From 619c577d3dc1e22b29540effbf9ae8e84c782481 Mon Sep 17 00:00:00 2001 From: Max Ghenis Date: Mon, 19 Oct 2026 08:39:22 -0400 Subject: [PATCH 1/2] Restore simulations with no persons or empty groups; refuse non-integer counts restore_simulation took a group's count as one past its highest member index. With no persons that was max() of an empty array, which raised; with groups that have no members last, it gave too small a count, and restoring their values then raised. It also took the number of persons from a group, so a system with no group entities raised UnboundLocalError. dump_simulation now writes each population's count (count.npy next to id.npy), and restore_simulation uses it. A dump without it gives the smallest count its ids and memberships allow. For a group, that differs from the old count only when it has more ids than one past its highest member index (its last groups have no members); for persons, only when the number of person ids or another group's memberships exceeds the last group's memberships. build_default_simulation used np.arange(count) after #590, which accepts 2.5, 2.0 and other non-integers; the simulation then failed in sum or calculate. It now applies operator.index first, which refuses exactly what range(count) refused before #590, and keeps int64 ids for every integer type. Co-Authored-By: Claude Opus 5.5 --- changelog.d/fix-empty-dump-restore.fixed.md | 1 + .../simulations/simulation_builder.py | 11 +- policyengine_core/tools/simulation_dumper.py | 55 +++-- .../test_default_simulation_count_property.py | 104 +++++++++ tests/core/test_dump_restore.py | 189 ++++++++++++++- tests/core/test_dump_restore_property.py | 220 ++++++++++++++++++ tests/core/test_group_population_empty.py | 21 +- tests/core/test_simulation_builder.py | 30 +++ 8 files changed, 615 insertions(+), 16 deletions(-) create mode 100644 changelog.d/fix-empty-dump-restore.fixed.md create mode 100644 tests/core/test_default_simulation_count_property.py create mode 100644 tests/core/test_dump_restore_property.py diff --git a/changelog.d/fix-empty-dump-restore.fixed.md b/changelog.d/fix-empty-dump-restore.fixed.md new file mode 100644 index 000000000..ce287c696 --- /dev/null +++ b/changelog.d/fix-empty-dump-restore.fixed.md @@ -0,0 +1 @@ +Restore dumped simulations that have no persons, or groups with no members, by recording each population's count in the dump; `build_default_simulation` again refuses a count that is not an integer when building. diff --git a/policyengine_core/simulations/simulation_builder.py b/policyengine_core/simulations/simulation_builder.py index 3e90a754c..c16a767d3 100644 --- a/policyengine_core/simulations/simulation_builder.py +++ b/policyengine_core/simulations/simulation_builder.py @@ -1,4 +1,5 @@ import copy +import operator import typing from typing import TYPE_CHECKING, Any, List @@ -334,8 +335,16 @@ def build_default_simulation( - There are ``count`` persons - There are ``count`` instances of each group entity, containing one person - Every person has, in each entity, the first role + + Raises: + TypeError: if ``count`` is not an integer, as ``range(count)`` + would. """ + # ``np.arange`` accepts 2.5 or 2.0, and the simulation then fails in + # ``sum`` or ``calculate``; ``operator.index`` refuses what ``range`` + # refuses. + id_count = operator.index(count) if simulation is None: simulation = Simulation( tax_benefit_system=tax_benefit_system, @@ -343,7 +352,7 @@ def build_default_simulation( ) for population in simulation.populations.values(): population.count = count - population.ids = np.arange(count, dtype=np.int64) + population.ids = np.arange(id_count, dtype=np.int64) if not population.entity.is_person: population.members_entity_id = ( population.ids diff --git a/policyengine_core/tools/simulation_dumper.py b/policyengine_core/tools/simulation_dumper.py index d351bb5d0..460de36b6 100644 --- a/policyengine_core/tools/simulation_dumper.py +++ b/policyengine_core/tools/simulation_dumper.py @@ -27,13 +27,21 @@ DUMP_FORMAT_FILE = "dump_format.txt" DUMP_FORMAT_VERSION = 1 +# In each population's directory in ``__entities__``: its count, which its +# ids and memberships do not always give (a group with no members holds no +# membership, and some simulations never set ids). ``restore_simulation`` +# infers the count of a population whose dump lacks this file (see +# ``_restore_count``); earlier versions ignore the file. +COUNT_FILE = "count.npy" + def dump_simulation(simulation, directory): """ Write simulation data to directory, so that it can be restored later. - Writes the simulation's branch values, including inherited values, and - lists the periods it calculated (see ``restore_simulation``). + Writes each population's ids and count and each group's memberships, + positions and roles; the simulation's branch values, including inherited + values; and the periods it calculated (see ``restore_simulation``). """ parent_directory = os.path.abspath(os.path.join(directory, os.pardir)) if not os.path.isdir(parent_directory): # To deal with reforms @@ -81,6 +89,11 @@ def restore_simulation(directory, tax_benefit_system, **kwargs): format marker that lists none may be one of those, so restoring it warns: ``apply_reform`` then keeps any calculated value in it as dumped instead of recalculating it. + + Each population gets the count the dump records, which may be zero. A + dump written before counts were recorded gives each count as the + smallest its ids and memberships allow: that is the count when the ids + are set, one per instance, as ``SimulationBuilder`` sets them. """ simulation = Simulation( tax_benefit_system, tax_benefit_system.instantiate_entities() @@ -88,15 +101,14 @@ def restore_simulation(directory, tax_benefit_system, **kwargs): entities_dump_dir = os.path.join(directory, "__entities__") for population in simulation.populations.values(): - if population.entity.is_person: - continue - person_count = _restore_entity(population, entities_dump_dir) - - for population in simulation.populations.values(): - if not population.entity.is_person: - continue _restore_entity(population, entities_dump_dir) - population.count = person_count + groups = [ + population + for population in simulation.populations.values() + if not population.entity.is_person + ] + for population in simulation.populations.values(): + _restore_count(population, entities_dump_dir, groups) variables_to_restore = [ variable for variable in os.listdir(directory) if variable != "__entities__" @@ -160,6 +172,7 @@ def _dump_entity(population, directory): path = os.path.join(directory, population.entity.key) os.mkdir(path) np.save(os.path.join(path, "id.npy"), population.ids) + np.save(os.path.join(path, COUNT_FILE), population.count) if population.entity.is_person: return @@ -202,9 +215,25 @@ def _restore_entity(population, directory): [encoded_roles == role.key for role in flattened_roles], [role for role in flattened_roles], ) - person_count = len(population.members_entity_id) - population.count = max(population.members_entity_id) + 1 - return person_count + + +def _restore_count(population, directory, groups): + """Set ``population.count`` from its dump, after every structure is restored. + + A dump without ``COUNT_FILE`` gives the smallest count its arrays allow: + one per id and, for a group, one past its highest member index (none + with no members); for persons, as many as any group has members. + """ + count_path = os.path.join(directory, population.entity.key, COUNT_FILE) + if os.path.exists(count_path): + population.count = np.load(count_path).item() + return + counts = [len(population.ids)] + if population.entity.is_person: + counts += [len(group.members_entity_id) for group in groups] + elif len(population.members_entity_id): + counts.append(int(np.max(population.members_entity_id)) + 1) + population.count = max(counts) def _restore_holder( diff --git a/tests/core/test_default_simulation_count_property.py b/tests/core/test_default_simulation_count_property.py new file mode 100644 index 000000000..e00d91716 --- /dev/null +++ b/tests/core/test_default_simulation_count_property.py @@ -0,0 +1,104 @@ +"""Property: ``build_default_simulation`` takes exactly the counts ``range`` takes. + +For any count: a Python integer (negative ones included), a NumPy integer of +any width or signedness, a 0-d integer array, a bool, a Python or NumPy float +(whole or not, NaN and infinities included), a 0-d float array, a NumPy bool, +a string, ``None``, a ``Decimal`` or a ``Fraction``: + +* building raises ``TypeError`` exactly when ``range(count)`` does, and then + leaves the simulation it was given as it was; +* otherwise every population's ids, and every group's memberships, are + ``range(count)`` as ``int64``. + +Examples are in ``test_simulation_builder.py`` and +``test_group_population_empty.py``. +""" + +from __future__ import annotations + +from decimal import Decimal +from fractions import Fraction + +import numpy as np +import pytest + +# The smoke job installs Core without the dev extra but collects every module. +pytest.importorskip("hypothesis") + +from hypothesis import given, settings # noqa: E402 +from hypothesis import strategies as st # noqa: E402 + +from policyengine_core.simulations import Simulation, SimulationBuilder # noqa: E402 +from policyengine_core.taxbenefitsystems import TaxBenefitSystem # noqa: E402 +from tests.core.test_group_population_empty import ( # noqa: E402 + HOUSEHOLD, + PERSON, +) + +SYSTEM = TaxBenefitSystem([PERSON, HOUSEHOLD]) +NUMPY_INTEGER_TYPES = [ + np.int8, + np.int16, + np.int32, + np.int64, + np.uint8, + np.uint16, + np.uint32, + np.uint64, + np.intp, + np.uintp, +] +_sizes = st.integers(0, 40) +_counts = st.one_of( + st.integers(-3, 40), + st.tuples(st.sampled_from(NUMPY_INTEGER_TYPES), _sizes).map( + lambda draw: draw[0](draw[1]) + ), + _sizes.map(np.array), + st.booleans(), + # Bounded, so a regression fails instead of allocating ``np.arange(1e300)``. + st.floats(-3, 40), + st.sampled_from([float("nan"), float("inf"), float("-inf")]), + _sizes.map(float), + st.floats(-3, 40).map(np.float64), + st.floats(-3, 40, width=32).map(np.float32), + st.floats(-3, 40).map(np.array), + st.booleans().map(np.bool_), + _sizes.map(str), + st.none(), + _sizes.map(Decimal), + _sizes.map(Fraction), +) + + +@settings(max_examples=300, deadline=None) +@given(_counts) +def test_default_simulation_takes_the_counts_range_takes(count): + try: + expected_ids = list(range(count)) + except TypeError: + expected_ids = None + simulation = Simulation( + tax_benefit_system=SYSTEM, populations=SYSTEM.instantiate_entities() + ) + + if expected_ids is None: + with pytest.raises(TypeError): + SimulationBuilder().build_default_simulation( + SYSTEM, count, simulation=simulation + ) + for population in simulation.populations.values(): + assert population.count == 0 + assert len(population.ids) == 0 + return + + built = SimulationBuilder().build_default_simulation( + SYSTEM, count, simulation=simulation + ) + for population in built.populations.values(): + assert population.count == count + assert population.ids.dtype == np.int64 + assert population.ids.tolist() == expected_ids + if not population.entity.is_person: + assert population.members_entity_id.dtype == np.int64 + assert population.members_entity_id.tolist() == expected_ids diff --git a/tests/core/test_dump_restore.py b/tests/core/test_dump_restore.py index 5282cb72d..82d2a8656 100644 --- a/tests/core/test_dump_restore.py +++ b/tests/core/test_dump_restore.py @@ -1,10 +1,13 @@ import shutil import tempfile +import numpy as np from numpy import testing from policyengine_core.country_template import situation_examples -from policyengine_core.simulations import SimulationBuilder +from policyengine_core.entities import build_entity +from policyengine_core.simulations import Simulation, SimulationBuilder +from policyengine_core.taxbenefitsystems import TaxBenefitSystem from policyengine_core.tools import simulation_dumper @@ -44,3 +47,187 @@ def test_dump(tax_benefit_system): testing.assert_array_equal(cached_value, calculated_value) shutil.rmtree(directory) + + +# A population's count is not always given by its ids and memberships: a +# group with no members holds no membership, and a simulation with no persons +# holds none at all. ``restore_simulation`` used to take a group's count as +# one past its highest member index, so it raised on ``max()`` of no members +# and gave too small a count when the last groups had none. + +EMPTY_GROUPS = { + "persons": { + "a": {"salary": {"2018-01": 1000}}, + "b": {"age": {"2018-01": 10}}, + "c": {}, + }, + # Households with no members first, in between and last. + "households": { + "h0": {}, + "h1": {"parents": ["a"], "children": ["b"], "rent": {"2018-01": 500}}, + "h2": {}, + "h3": {"parents": ["c"]}, + "h4": {"rent": {"2018-01": 700}}, + }, +} + + +def _round_trip(simulation, tax_benefit_system, directory, *, without_counts=False): + simulation_dumper.dump_simulation(simulation, str(directory)) + if without_counts: + # As dumps written before counts were recorded. + for path in directory.glob(f"__entities__/*/{simulation_dumper.COUNT_FILE}"): + path.unlink() + return simulation_dumper.restore_simulation(str(directory), tax_benefit_system) + + +def _assert_same_structure(simulation, restored): + for key, population in simulation.populations.items(): + restored_population = restored.populations[key] + assert restored_population.count == population.count, key + testing.assert_array_equal(restored_population.ids, population.ids) + if population.entity.is_person: + continue + for name in ("members_entity_id", "members_position", "members_role"): + original = np.asarray(getattr(population, name)) + restored_array = np.asarray(getattr(restored_population, name)) + testing.assert_array_equal(restored_array, original, err_msg=name) + if original.dtype.kind in "iu": + assert restored_array.dtype.kind in "iu", name + + +def test_dump_restore_with_no_persons(tax_benefit_system, tmp_path): + simulation = SimulationBuilder().build_default_simulation( + tax_benefit_system, count=0 + ) + simulation.set_input("salary", "2018-01", np.array([])) + simulation.set_input("rent", "2018-01", np.array([])) + calculated = simulation.calculate("disposable_income", "2018-01") + + restored = _round_trip(simulation, tax_benefit_system, tmp_path / "dump") + + _assert_same_structure(simulation, restored) + assert [population.count for population in restored.populations.values()] == [ + 0, + 0, + ] + holder = restored.person.get_holder("disposable_income") + testing.assert_array_equal(holder.get_array("2018-01"), calculated) + assert holder.is_derived("2018-01") + assert restored.calculate("household_income", "2018-01").tolist() == [] + assert restored.household.max(np.array([])).tolist() == [] + + +def test_dump_restore_with_empty_groups_among_populated_ones( + tax_benefit_system, tmp_path +): + simulation = SimulationBuilder().build_from_entities( + tax_benefit_system, EMPTY_GROUPS + ) + household_income = simulation.calculate("household_income", "2018-01") + assert simulation.household.count == 5 + assert simulation.household.members_entity_id.tolist() == [1, 1, 3] + + restored = _round_trip(simulation, tax_benefit_system, tmp_path / "dump") + + _assert_same_structure(simulation, restored) + testing.assert_array_equal( + restored.household.get_holder("household_income").get_array("2018-01"), + household_income, + ) + testing.assert_array_equal( + restored.calculate("rent", "2018-01"), [0, 500, 0, 0, 700] + ) + testing.assert_array_equal( + restored.household.nb_persons(), simulation.household.nb_persons() + ) + + +def test_restore_infers_the_counts_of_a_dump_without_them(tax_benefit_system, tmp_path): + # Every simulation built by ``SimulationBuilder`` has one id per instance. + simulation = SimulationBuilder().build_from_entities( + tax_benefit_system, EMPTY_GROUPS + ) + simulation.calculate("household_income", "2018-01") + + restored = _round_trip( + simulation, tax_benefit_system, tmp_path / "groups", without_counts=True + ) + + _assert_same_structure(simulation, restored) + testing.assert_array_equal( + restored.calculate("household_income", "2018-01"), + simulation.calculate("household_income", "2018-01"), + ) + + no_one = SimulationBuilder().build_default_simulation(tax_benefit_system, count=0) + restored = _round_trip( + no_one, tax_benefit_system, tmp_path / "no_one", without_counts=True + ) + _assert_same_structure(no_one, restored) + + +def test_restore_infers_counts_from_memberships_without_group_ids( + tax_benefit_system, tmp_path +): + # As ``scripts/simulation_generator.py`` builds one: counts and + # memberships, but no group ids. Each group has members. + simulation = Simulation( + tax_benefit_system, tax_benefit_system.instantiate_entities() + ) + simulation.persons.ids = np.arange(4) + simulation.persons.count = 4 + simulation.household.count = 3 + simulation.household.members_entity_id = np.array([0, 1, 1, 2]) + simulation.set_input("rent", "2018-01", np.array([1.0, 2.0, 3.0])) + assert len(simulation.household.ids) == 0 + + restored = _round_trip( + simulation, tax_benefit_system, tmp_path / "dump", without_counts=True + ) + + assert restored.person.count == 4 + assert restored.household.count == 3 + testing.assert_array_equal( + restored.household.members_entity_id, simulation.household.members_entity_id + ) + testing.assert_array_equal(restored.calculate("rent", "2018-01"), [1, 2, 3]) + + +def test_dump_keeps_a_count_its_ids_and_memberships_do_not_give( + tax_benefit_system, tmp_path +): + # No group ids, and the last household has no members: only the dumped + # count says there are three households. + simulation = Simulation( + tax_benefit_system, tax_benefit_system.instantiate_entities() + ) + simulation.persons.ids = np.arange(2) + simulation.persons.count = 2 + simulation.household.count = 3 + simulation.household.members_entity_id = np.array([0, 1]) + simulation.set_input("rent", "2018-01", np.array([1.0, 2.0, 3.0])) + + restored = _round_trip(simulation, tax_benefit_system, tmp_path / "dump") + + assert restored.household.count == 3 + testing.assert_array_equal(restored.calculate("rent", "2018-01"), [1, 2, 3]) + + +def test_dump_restore_without_group_entities(tmp_path): + # With no group to read the number of persons from, restoring raised + # UnboundLocalError. + system = TaxBenefitSystem( + [build_entity("person", "persons", "Person", is_person=True)] + ) + for count in (0, 3): + simulation = SimulationBuilder().build_default_simulation(system, count=count) + for without_counts in (False, True): + restored = _round_trip( + simulation, + system, + tmp_path / f"dump-{count}-{without_counts}", + without_counts=without_counts, + ) + assert restored.persons.count == count + testing.assert_array_equal(restored.persons.ids, simulation.persons.ids) diff --git a/tests/core/test_dump_restore_property.py b/tests/core/test_dump_restore_property.py new file mode 100644 index 000000000..197d33b8f --- /dev/null +++ b/tests/core/test_dump_restore_property.py @@ -0,0 +1,220 @@ +"""Property: a dumped simulation restores with the same structure and values. + +For any number of persons (0 to 6) and of households (0 to 4, at least one +when there are persons), any assignment of persons to households (so a +household may have no members, first, last or in between, and there may be +no one at all), any roles, ids given as integers or strings or never set (as +``scripts/simulation_generator.py`` leaves them for groups), any inputs of +each value type the country template has (float, int, date, Enum and str, +monthly and ETERNITY), and any calculations before the dump: + +1. **Structure.** Every population's count and ids, and every group's + memberships, positions and roles, are restored, including zero counts. +2. **Values.** Every holder stores the same values, byte for byte, with the + same derived marks. +3. **Dumps without counts.** With the count files removed, as in dumps + written before they were, restoring gives the same structure and values + whenever the ids and memberships give the counts: when the ids are set, + one per instance, as ``SimulationBuilder`` sets them, or when the last + household has members. (Otherwise only the count file gives the count.) +4. **Same answers.** Any calculation on the restored simulation gives what + it gives on the dumped one, including any error. + +Examples are in ``test_dump_restore.py``. +""" + +from __future__ import annotations + +import datetime +import pathlib +import tempfile + +import numpy as np +import pytest + +# The smoke job installs Core without the dev extra but collects every module. +pytest.importorskip("hypothesis") + +from hypothesis import HealthCheck, assume, given, settings # noqa: E402 +from hypothesis import strategies as st # noqa: E402 + +from policyengine_core.country_template import CountryTaxBenefitSystem # noqa: E402 +from policyengine_core.enums import EnumArray # noqa: E402 +from policyengine_core.simulations import SimulationBuilder # noqa: E402 +from policyengine_core.tools.simulation_dumper import ( # noqa: E402 + COUNT_FILE, + dump_simulation, + restore_simulation, +) + +SYSTEM = CountryTaxBenefitSystem() +ROLES = SYSTEM.entities_by_singular()["household"].flattened_roles +MONTHS = ["2018-01", "2018-02"] + +_floats = st.floats(-1e5, 1e5, allow_nan=False, width=32) +_dates = st.dates(datetime.date(1900, 1, 1), datetime.date(2030, 12, 31)).map( + np.datetime64 +) +_statuses = st.sampled_from(["owner", "tenant", "free_lodger", "homeless"]) +# Variable, entity key, periods, element strategy and array dtype. +INPUTS = [ + ("salary", "person", MONTHS, _floats, float), + ("age", "person", MONTHS, st.integers(0, 120), int), + ("birth", "person", ["ETERNITY"], _dates, "datetime64[D]"), + ("rent", "household", MONTHS, _floats, float), + ("housing_occupancy_status", "household", MONTHS, _statuses, str), + ("postal_code", "household", MONTHS, st.text("ab12", max_size=5), str), +] +CALCULATIONS = [ + "disposable_income", + "basic_income", + "income_tax", + "household_income", + "housing_allowance", + "parenting_allowance", + "housing_tax", +] + + +@st.composite +def _simulations(draw): + person_count = draw(st.integers(0, 6)) + household_count = draw(st.integers(1 if person_count else 0, 4)) + memberships = draw( + st.lists( + st.integers(0, max(household_count - 1, 0)), + min_size=person_count, + max_size=person_count, + ) + ) + roles = draw( + st.lists(st.sampled_from(ROLES), min_size=person_count, max_size=person_count) + ) + id_kind = draw(st.sampled_from(["int", "str", "unset"])) + counts = {"person": person_count, "household": household_count} + inputs = [] + for _ in range(draw(st.integers(0, 6))): + name, entity, input_periods, elements, dtype = draw(st.sampled_from(INPUTS)) + values = draw( + st.lists(elements, min_size=counts[entity], max_size=counts[entity]) + ) + inputs.append( + (name, draw(st.sampled_from(input_periods)), np.array(values, dtype=dtype)) + ) + calculations = draw( + st.lists(st.tuples(st.sampled_from(CALCULATIONS), st.sampled_from(MONTHS))) + ) + return ( + person_count, + household_count, + memberships, + roles, + id_kind, + inputs, + calculations, + ) + + +def _build(person_count, household_count, memberships, roles, id_kind): + """A simulation with this structure, set as ``SimulationBuilder`` sets one.""" + simulation = SimulationBuilder().build_default_simulation( + SYSTEM, count=person_count + ) + + def ids(prefix, count): + if id_kind == "str": + return np.array([f"{prefix}{index}" for index in range(count)], dtype=str) + if id_kind == "unset": + return [] + return np.arange(count, dtype=np.int64) + + simulation.persons.ids = ids("p", person_count) + households = simulation.household + households.count = household_count + households.ids = ids("h", household_count) + households.members_entity_id = np.array(memberships, dtype=np.int64) + households.members_role = np.array(roles, dtype=object) + return simulation + + +def _result(function): + try: + value = function() + except Exception as error: # Compare failures as well as values. + return ("error", type(error).__name__, str(error)) + if isinstance(value, EnumArray): + value = value.view(np.ndarray) + value = np.asarray(value) + return ("array", value.dtype.str, value.shape, value.tobytes()) + + +def _structure(simulation): + structure = {} + for key, population in simulation.populations.items(): + ids = np.asarray(population.ids) + entry = {"count": population.count, "ids": (ids.dtype.str, ids.tolist())} + if not population.entity.is_person: + for name in ("members_entity_id", "members_position"): + array = np.asarray(getattr(population, name)) + entry[name] = (array.dtype.kind in "iu", array.tolist()) + entry["members_role"] = [role.key for role in population.members_role] + structure[key] = entry + return structure + + +def _stored(simulation): + """Every value and its derived mark, as bytes, keyed by variable and period.""" + stored = {} + for population in simulation.populations.values(): + for name, holder in population._holders.items(): + for period in holder.get_known_periods(): + value = holder.get_array(period) + if isinstance(value, EnumArray): + value = value.view(np.ndarray) + stored[(name, str(period))] = ( + value.dtype.str, + value.tobytes(), + holder.is_derived(period), + ) + return stored + + +@settings( + max_examples=200, + deadline=None, + suppress_health_check=[HealthCheck.too_slow, HealthCheck.data_too_large], +) +@given(case=_simulations(), without_counts=st.booleans()) +def test_dump_then_restore_keeps_structure_and_values(case, without_counts): + *structure, inputs, calculations = case + person_count, household_count, memberships, _, id_kind = structure + last_household_empty = ( + household_count > 0 and household_count - 1 not in memberships + ) + assume(not (without_counts and id_kind == "unset" and last_household_empty)) + dumped = _build(*structure) + for name, period, values in inputs: + dumped.set_input(name, period, values) + for request in calculations: + _result(lambda: dumped.calculate(*request)) + + with tempfile.TemporaryDirectory(prefix="core-dump-") as directory: + dump_directory = pathlib.Path(directory) / "dump" + dump_simulation(dumped, str(dump_directory)) + count_files = list(dump_directory.glob(f"__entities__/*/{COUNT_FILE}")) + assert len(count_files) == len(dumped.populations) + if without_counts: + for path in count_files: + path.unlink() + restored = restore_simulation(str(dump_directory), SYSTEM) + + # 1. Structure, and 3. without counts. + assert _structure(restored) == _structure(dumped) + # 2. Values. + assert _stored(restored) == _stored(dumped) + # 4. Same answers. + for name in CALCULATIONS: + for period in MONTHS: + assert _result(lambda: restored.calculate(name, period)) == _result( + lambda: dumped.calculate(name, period) + ), (name, period) diff --git a/tests/core/test_group_population_empty.py b/tests/core/test_group_population_empty.py index 7e7576769..b0b4bf6f0 100644 --- a/tests/core/test_group_population_empty.py +++ b/tests/core/test_group_population_empty.py @@ -84,7 +84,26 @@ def test_default_simulation_of_no_one(): assert households.all(no_values > 0).tolist() == [] -@pytest.mark.parametrize("count_type", [int, np.int32, np.int64, np.uint64]) +INTEGER_COUNT_TYPES = [ + int, + np.int8, + np.int16, + np.int32, + np.int64, + np.uint8, + np.uint16, + np.uint32, + np.uint64, + np.intp, + np.uintp, + np.intc, + np.longlong, + # A 0-d integer array, which ``range`` also takes. + np.array, +] + + +@pytest.mark.parametrize("count_type", INTEGER_COUNT_TYPES) @pytest.mark.parametrize("count_value", [0, 1, 4]) def test_default_simulation_keeps_integer_ids(count_type, count_value): count = count_type(count_value) diff --git a/tests/core/test_simulation_builder.py b/tests/core/test_simulation_builder.py index f838d6b3c..049bedf78 100644 --- a/tests/core/test_simulation_builder.py +++ b/tests/core/test_simulation_builder.py @@ -1,6 +1,9 @@ import datetime +from decimal import Decimal +from fractions import Fraction from typing import Iterable +import numpy as np import pytest from policyengine_core import periods, tools @@ -85,6 +88,33 @@ def test_build_default_simulation(tax_benefit_system): ).all() +@pytest.mark.parametrize( + "count", + [ + 2.5, + 2.0, + np.float64(2.0), + np.float32(1.0), + np.array(2.0), + np.True_, + "2", + None, + Decimal(2), + Fraction(2), + ], + ids=repr, +) +def test_build_default_simulation_refuses_a_count_that_is_not_an_integer( + tax_benefit_system, count +): + # ``range(count)`` refused these when building. ``np.arange`` took the + # floats, and the simulation then failed in ``sum`` or ``calculate``. + with pytest.raises(TypeError): + range(count) + with pytest.raises(TypeError, match="cannot be interpreted as an integer|index"): + SimulationBuilder().build_default_simulation(tax_benefit_system, count) + + def test_explicit_singular_entities(tax_benefit_system): assert SimulationBuilder().explicit_singular_entities( tax_benefit_system, From 44e6914f56fb57e1b198e0cd2f60f11a51b047dc Mon Sep 17 00:00:00 2001 From: Max Ghenis Date: Mon, 19 Oct 2026 08:56:23 -0400 Subject: [PATCH 2/2] Dump each count as int64, and none for a count that is not an integer Review of 619c577: np.save(population.count) saved the count object as given. A masked-array count made the dump raise NotImplementedError, and an integer-like object with only __index__ wrote an object array that restoring refuses (allow_pickle=False). The dump now writes np.int64(operator.index(count)). A count that is not an integer, which only code that sets it by hand can make, gets no count file, so restoring infers it as for an older dump, and the dump does not raise where it used to work. The round-trip property also compares each value's array type, shape and enum class, and checks that every count file is int64. Co-Authored-By: Claude Opus 5.5 --- policyengine_core/tools/simulation_dumper.py | 22 +++++--- tests/core/test_dump_restore.py | 54 ++++++++++++++++++++ tests/core/test_dump_restore_property.py | 13 ++++- 3 files changed, 82 insertions(+), 7 deletions(-) diff --git a/policyengine_core/tools/simulation_dumper.py b/policyengine_core/tools/simulation_dumper.py index 460de36b6..ceb726db4 100644 --- a/policyengine_core/tools/simulation_dumper.py +++ b/policyengine_core/tools/simulation_dumper.py @@ -1,6 +1,7 @@ # -*- coding: utf-8 -*- +import operator import os import warnings @@ -27,11 +28,12 @@ DUMP_FORMAT_FILE = "dump_format.txt" DUMP_FORMAT_VERSION = 1 -# In each population's directory in ``__entities__``: its count, which its -# ids and memberships do not always give (a group with no members holds no -# membership, and some simulations never set ids). ``restore_simulation`` -# infers the count of a population whose dump lacks this file (see -# ``_restore_count``); earlier versions ignore the file. +# In each population's directory in ``__entities__``: its count as an int64 +# scalar, which its ids and memberships do not always give (a group with no +# members holds no membership, and some simulations never set ids). +# ``restore_simulation`` infers the count of a population whose dump lacks +# this file (see ``_restore_count``), as for a count that is not an integer; +# earlier versions ignore the file. COUNT_FILE = "count.npy" @@ -172,7 +174,15 @@ def _dump_entity(population, directory): path = os.path.join(directory, population.entity.key) os.mkdir(path) np.save(os.path.join(path, "id.npy"), population.ids) - np.save(os.path.join(path, COUNT_FILE), population.count) + try: + # Any integer, such as a NumPy one or a 0-d or masked array, as int64. + count = np.int64(operator.index(population.count)) + except TypeError: + # A count that is not an integer can only be set by hand: restore + # infers it, as it does for an older dump. + count = None + if count is not None: + np.save(os.path.join(path, COUNT_FILE), count) if population.entity.is_person: return diff --git a/tests/core/test_dump_restore.py b/tests/core/test_dump_restore.py index 82d2a8656..112a880c2 100644 --- a/tests/core/test_dump_restore.py +++ b/tests/core/test_dump_restore.py @@ -2,6 +2,7 @@ import tempfile import numpy as np +import pytest from numpy import testing from policyengine_core.country_template import situation_examples @@ -214,6 +215,59 @@ def test_dump_keeps_a_count_its_ids_and_memberships_do_not_give( testing.assert_array_equal(restored.calculate("rent", "2018-01"), [1, 2, 3]) +class _IntegerLike: + """An integer only through ``__index__``, as ``range`` takes it.""" + + def __init__(self, value): + self.value = value + + def __index__(self): + return self.value + + +@pytest.mark.parametrize( + "count", + [np.uint8(2), np.array(2), np.ma.array(2, dtype=np.int64), _IntegerLike(2)], + ids=["uint8", "0-d array", "masked array", "__index__"], +) +def test_dump_writes_any_integer_count_as_int64(tax_benefit_system, tmp_path, count): + # ``np.save`` of the count itself failed for a masked array and wrote an + # object array, which restoring refuses, for an ``__index__`` object. + simulation = SimulationBuilder().build_default_simulation( + tax_benefit_system, count=count + ) + + restored = _round_trip(simulation, tax_benefit_system, tmp_path / "dump") + + count_files = list( + (tmp_path / "dump").glob(f"__entities__/*/{simulation_dumper.COUNT_FILE}") + ) + assert len(count_files) == 2 + for path in count_files: + assert np.load(path).dtype == np.int64 + for population in restored.populations.values(): + assert type(population.count) is int and population.count == 2 + testing.assert_array_equal(population.ids, [0, 1]) + + +def test_dump_of_a_count_that_is_not_an_integer_infers_it(tax_benefit_system, tmp_path): + # Only set by hand: the dump records no count, and restoring infers it + # from the ids, as for an older dump. + simulation = SimulationBuilder().build_default_simulation( + tax_benefit_system, count=2 + ) + simulation.household.count = 2.0 + simulation.set_input("rent", "2018-01", np.array([1.0, 2.0])) + + restored = _round_trip(simulation, tax_benefit_system, tmp_path / "dump") + + assert not ( + tmp_path / "dump" / "__entities__" / "household" / simulation_dumper.COUNT_FILE + ).exists() + assert restored.household.count == 2 + testing.assert_array_equal(restored.calculate("rent", "2018-01"), [1, 2]) + + def test_dump_restore_without_group_entities(tmp_path): # With no group to read the number of persons from, restoring raised # UnboundLocalError. diff --git a/tests/core/test_dump_restore_property.py b/tests/core/test_dump_restore_property.py index 197d33b8f..9d1e183f9 100644 --- a/tests/core/test_dump_restore_property.py +++ b/tests/core/test_dump_restore_property.py @@ -163,16 +163,25 @@ def _structure(simulation): def _stored(simulation): - """Every value and its derived mark, as bytes, keyed by variable and period.""" + """Every value and its derived mark, as bytes, keyed by variable and period. + + An Enum value also keeps its enum class. + """ stored = {} for population in simulation.populations.values(): for name, holder in population._holders.items(): for period in holder.get_known_periods(): value = holder.get_array(period) + array_type = type(value).__name__ + possible_values = None if isinstance(value, EnumArray): + possible_values = value.possible_values value = value.view(np.ndarray) stored[(name, str(period))] = ( + array_type, + possible_values, value.dtype.str, + value.shape, value.tobytes(), holder.is_derived(period), ) @@ -203,6 +212,8 @@ def test_dump_then_restore_keeps_structure_and_values(case, without_counts): dump_simulation(dumped, str(dump_directory)) count_files = list(dump_directory.glob(f"__entities__/*/{COUNT_FILE}")) assert len(count_files) == len(dumped.populations) + for path in count_files: + assert np.load(path).dtype == np.int64 if without_counts: for path in count_files: path.unlink()