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..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,13 +28,22 @@ DUMP_FORMAT_FILE = "dump_format.txt" DUMP_FORMAT_VERSION = 1 +# 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" + 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 +91,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 +103,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 +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) + 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 @@ -202,9 +225,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..112a880c2 100644 --- a/tests/core/test_dump_restore.py +++ b/tests/core/test_dump_restore.py @@ -1,10 +1,14 @@ import shutil import tempfile +import numpy as np +import pytest 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 +48,240 @@ 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]) + + +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. + 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..9d1e183f9 --- /dev/null +++ b/tests/core/test_dump_restore_property.py @@ -0,0 +1,231 @@ +"""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. + + 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), + ) + 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) + for path in count_files: + assert np.load(path).dtype == np.int64 + 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,