Restore simulations with no persons or empty groups, and refuse non-integer default counts

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.

policyengine_core/simulations/simulation_builder.py

@@ -1,4 +1,5 @@
 import copy
+import operator
 import typing
 from typing import TYPE_CHECKING, Any, List
 
@@ -334,16 +335,24 @@ 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,
                 populations=tax_benefit_system.instantiate_entities(),
             )
         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

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,22 +91,26 @@ 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()
     )
 
     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(

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