From d647f78d51fad3225562f638056270a7ab1ae8d6 Mon Sep 17 00:00:00 2001 From: Max Ghenis Date: Tue, 6 Oct 2026 13:37:01 -0400 Subject: [PATCH 1/4] Join each person to the group their ID names SimulationBuilder.join_with_persons ranked the distinct group IDs persons belong to and used that rank to index the declared IDs in sorted order. A group with no members has no rank, so when one sorted before an occupied group every later person moved down one group. With groups ["a", "b"] and one person in "b" (income 100), members_entity_id was [0] and the household sums [100, 0]. Simulation.build_from_dataset calls it for every group entity. join_with_persons now looks each person's group ID up among all declared IDs (group_positions: stable argsort plus searchsorted). Empty groups can sit anywhere in the declared order. A person whose ID names no declared group, and a declared ID that repeats, raise ValueError instead of silently landing somewhere. The flat-file path in build_from_dataset relied on the old ranking: it declared groups 0..n-1 and passed the raw membership IDs. It now numbers the distinct membership IDs itself, which gives the same mapping. A flat file with an unsuffixed household_id column skipped that path and declared the previous entity's IDs as its households (three persons in two households made three); it now takes the same path. Adds regressions (the audit witness, an interior empty group, an ARRAYS dataset with an empty household, both flat-file paths, the two errors) and a Hypothesis property comparing membership, sums, counts and projections with a dict oracle over random IDs, orders, empty groups and assignments. Co-Authored-By: Claude Opus 5.5 --- changelog.d/fix-join-group-identity.fixed.md | 1 + policyengine_core/simulations/simulation.py | 29 ++-- .../simulations/simulation_builder.py | 56 ++++++- tests/core/test_join_with_persons_identity.py | 148 ++++++++++++++++++ ...est_join_with_persons_identity_property.py | 91 +++++++++++ 5 files changed, 302 insertions(+), 23 deletions(-) create mode 100644 changelog.d/fix-join-group-identity.fixed.md create mode 100644 tests/core/test_join_with_persons_identity.py create mode 100644 tests/core/test_join_with_persons_identity_property.py diff --git a/changelog.d/fix-join-group-identity.fixed.md b/changelog.d/fix-join-group-identity.fixed.md new file mode 100644 index 000000000..45f7439da --- /dev/null +++ b/changelog.d/fix-join-group-identity.fixed.md @@ -0,0 +1 @@ +Persons joined to groups from a dataset now land in the group their ID names even when a group with no members sorts before theirs; an ID that names no declared group, or a repeated group ID, now raises instead of silently picking a group. A flat-file dataset with a `household_id`-style column no longer declares the previous entity's IDs as its groups. diff --git a/policyengine_core/simulations/simulation.py b/policyengine_core/simulations/simulation.py index 1efe0c6b5..8e9ea4f39 100644 --- a/policyengine_core/simulations/simulation.py +++ b/policyengine_core/simulations/simulation.py @@ -574,30 +574,27 @@ def get_eternity_array(name): for group_entity in self.tax_benefit_system.group_entities: entity_id_field = f"{group_entity.key}_id" + person_membership_id_field = f"{person_entity.key}_{group_entity.key}_id" if self.dataset.data_format != Dataset.FLAT_FILE: assert entity_id_field in data, ( f"Missing {entity_id_field} column in the dataset. Each group entity must have an ID array defined for ETERNITY." ) entity_ids = get_eternity_array(entity_id_field) - elif entity_id_field not in data: - entity_id_field_values = get_eternity_array( - f"person_{group_entity.key}_id" - ) - if entity_id_field_values is not None: - entity_ids = np.arange(len(np.unique(entity_id_field_values))) - else: - entity_ids = np.arange(len(data[list(data.keys())[0]])) - - builder.declare_entity(group_entity.key, entity_ids) - - person_membership_id_field = f"{person_entity.key}_{group_entity.key}_id" - if self.dataset.data_format != Dataset.FLAT_FILE: assert person_membership_id_field in data, ( f"Missing {person_membership_id_field} column in the dataset. Each group entity must have a person membership array defined for ETERNITY." ) - elif person_membership_id_field not in data: - data[person_membership_id_field] = np.arange(len(data)) - person_membership_ids = get_eternity_array(person_membership_id_field) + person_membership_ids = get_eternity_array(person_membership_id_field) + else: + # A flat file has one row per person, so it declares each + # group only through its members' IDs. Number the distinct + # IDs 0, 1, ... and give each person their group's number. + membership_values = get_eternity_array(person_membership_id_field) + distinct_ids, person_membership_ids = np.unique( + membership_values, return_inverse=True + ) + entity_ids = np.arange(len(distinct_ids)) + + builder.declare_entity(group_entity.key, entity_ids) person_role_field = f"{person_entity.key}_{group_entity.key}_role" if person_role_field in data: diff --git a/policyengine_core/simulations/simulation_builder.py b/policyengine_core/simulations/simulation_builder.py index e3dee0cc5..b4bba8732 100644 --- a/policyengine_core/simulations/simulation_builder.py +++ b/policyengine_core/simulations/simulation_builder.py @@ -28,6 +28,50 @@ ) from policyengine_core.variables import Variable + + +def group_positions( + group_ids: ArrayLike, persons_group_ids: ArrayLike, entity_key: str = "group" +) -> np.ndarray: + """Return the position in ``group_ids`` of each person's group. + + ``members_entity_id`` indexes the arrays of a group population, which + follow the order ``group_ids`` were declared in. A person belongs to the + group whose declared ID equals theirs, wherever that group sits in the + declared order and whether or not other groups have members. + + Raises: + ValueError: if the declared IDs repeat, or a person belongs to an ID + that was not declared. + """ + group_ids = np.asarray(group_ids) + persons_group_ids = np.asarray(persons_group_ids) + sorter = np.argsort(group_ids, kind="stable") + sorted_ids = group_ids[sorter] + repeated = sorted_ids[1:][sorted_ids[1:] == sorted_ids[:-1]] + if len(repeated) > 0: + raise ValueError( + f"{entity_key} IDs must be unique, but these repeat: " + f"{np.unique(repeated)[:5].tolist()}." + ) + try: + positions = np.searchsorted(sorted_ids, persons_group_ids) + except TypeError as error: + raise ValueError( + f"Person {entity_key} IDs ({persons_group_ids.dtype}) cannot be " + f"compared with the declared {entity_key} IDs ({group_ids.dtype})." + ) from error + declared = positions < len(sorted_ids) + declared[declared] = sorted_ids[positions[declared]] == persons_group_ids[declared] + if not declared.all(): + undeclared = np.unique(persons_group_ids[~declared]) + raise ValueError( + f"{int((~declared).sum())} person(s) belong to {entity_key} IDs " + f"that were not declared: {undeclared[:5].tolist()}." + ) + return sorter[positions] + + from datetime import datetime @@ -281,13 +325,11 @@ def join_with_persons( persons_group_assignment: ArrayLike, roles: typing.Iterable[str], ) -> None: - # Maps group's identifiers to a 0-based integer range, for indexing into members_roles (see PR#876) - group_sorted_indices = np.unique(persons_group_assignment, return_inverse=True)[ - 1 - ] - group_population.members_entity_id = np.argsort(group_population.ids)[ - group_sorted_indices - ] + group_population.members_entity_id = group_positions( + group_population.ids, + persons_group_assignment, + group_population.entity.key, + ) flattened_roles = group_population.entity.flattened_roles roles_array = np.array(roles) diff --git a/tests/core/test_join_with_persons_identity.py b/tests/core/test_join_with_persons_identity.py new file mode 100644 index 000000000..f4f4d7223 --- /dev/null +++ b/tests/core/test_join_with_persons_identity.py @@ -0,0 +1,148 @@ +"""Joining persons to groups keeps each person in the group their ID names. + +``SimulationBuilder.join_with_persons`` turns each person's group ID into +``members_entity_id``, the position of that group among the declared IDs. +Empty groups may sit anywhere in the declared order (``test_bincount_minlength`` +covers one at the end), so the join must look each ID up among all declared +IDs, not rank it among the occupied ones. Ranking moved a person in group +``"b"`` into group ``"a"`` whenever ``"a"`` had no members. + +``test_join_with_persons_identity_property.py`` checks the same identity for +random IDs, declared orders, empty groups and assignments. +""" + +from __future__ import annotations + +import numpy as np +import pandas as pd +import pytest + +from policyengine_core.data import Dataset +from policyengine_core.simulations import Simulation, SimulationBuilder +from policyengine_core.simulations.simulation_builder import group_positions + + +def _join(tax_benefit_system, group_ids, persons_group_ids): + builder = SimulationBuilder() + builder.create_entities(tax_benefit_system) + builder.declare_person_entity("person", list(range(len(persons_group_ids)))) + household = builder.declare_entity("household", group_ids) + builder.join_with_persons( + household, persons_group_ids, ["first_parent"] * len(persons_group_ids) + ) + return household + + +def test_person_joins_named_group_when_earlier_group_is_empty(tax_benefit_system): + # The audit witness: group "a" is empty and sorts before group "b". + household = _join(tax_benefit_system, ["a", "b"], ["b"]) + + assert household.members_entity_id.tolist() == [1] + assert household.sum(np.array([100.0])).tolist() == [0.0, 100.0] + assert household.project(np.array([10.0, 20.0])).tolist() == [20.0] + + +def test_person_joins_named_group_around_an_interior_empty_group( + tax_benefit_system, +): + household = _join(tax_benefit_system, ["c", "a", "b"], ["c", "a", "a"]) + + assert household.members_entity_id.tolist() == [0, 1, 1] + assert household.nb_persons().tolist() == [1, 2, 0] + + +def test_join_rejects_an_undeclared_group(tax_benefit_system): + with pytest.raises(ValueError, match="not declared: \\['z'\\]"): + _join(tax_benefit_system, ["a", "b"], ["a", "z"]) + + +def test_join_rejects_repeated_group_ids(tax_benefit_system): + with pytest.raises(ValueError, match="must be unique"): + _join(tax_benefit_system, ["a", "b", "a"], ["a", "b"]) + + +def test_group_positions_of_no_persons_is_empty(): + assert group_positions(["a", "b"], []).tolist() == [] + assert group_positions([], []).tolist() == [] + + +def test_dataset_person_joins_named_household_when_another_is_empty( + tax_benefit_system, tmp_path +): + # Household 20 is declared but empty; the person in household 30 must + # see household 30's rent, not household 20's. + class InteriorEmptyHousehold(Dataset): + name = "interior_empty_household" + label = "Interior empty household" + file_path = tmp_path / "interior_empty_household.h5" + data_format = Dataset.ARRAYS + time_period = "2024" + + InteriorEmptyHousehold().save_dataset( + { + "person_id": np.array([1, 2]), + "household_id": np.array([10, 20, 30]), + "person_household_id": np.array([10, 30]), + # HDF5 stores no unicode strings, so give roles by index. + "person_household_role": np.array([0, 0]), + } + ) + simulation = Simulation( + tax_benefit_system=tax_benefit_system, + dataset=InteriorEmptyHousehold(), + ) + simulation.set_input("rent", "2024-01", np.array([100.0, 200.0, 300.0])) + + household = simulation.populations["household"] + assert household.members_entity_id.tolist() == [0, 2] + assert simulation.calculate("rent", "2024-01", map_to="person").tolist() == [ + 100.0, + 300.0, + ] + + +def test_flat_file_numbers_households_by_their_members(tax_benefit_system): + # A flat file has one row per person; households 500 and 700 become + # households 0 and 1, in order of their IDs. + dataframe = pd.DataFrame( + { + "person_id__2024": [10, 11, 12], + "person_household_id__2024": [700, 500, 700], + "salary__2024-01": [1.0, 2.0, 4.0], + } + ) + simulation = Simulation( + tax_benefit_system=tax_benefit_system, + dataset=Dataset.from_dataframe(dataframe, "2024"), + ) + + household = simulation.populations["household"] + assert household.ids.tolist() == [0, 1] + assert household.members_entity_id.tolist() == [1, 0, 1] + assert simulation.calculate("salary", "2024-01", map_to="household").tolist() == [ + 2.0, + 5.0, + ] + + +def test_flat_file_with_household_id_column_counts_its_households( + tax_benefit_system, +): + # A ``household_id`` column without a period suffix used to leave the + # previous entity's IDs declared: three persons made three households. + dataframe = pd.DataFrame( + { + "person_id": [10, 11, 12], + "household_id": [500, 500, 700], + "person_household_id": [500, 500, 700], + "salary__2024-01": [1.0, 2.0, 4.0], + } + ) + simulation = Simulation( + tax_benefit_system=tax_benefit_system, + dataset=Dataset.from_dataframe(dataframe, "2024"), + ) + + household = simulation.populations["household"] + assert household.count == 2 + assert household.members_entity_id.tolist() == [0, 0, 1] diff --git a/tests/core/test_join_with_persons_identity_property.py b/tests/core/test_join_with_persons_identity_property.py new file mode 100644 index 000000000..17c979cc2 --- /dev/null +++ b/tests/core/test_join_with_persons_identity_property.py @@ -0,0 +1,91 @@ +"""Property: joining persons to groups keeps each person in their named group. + +For any distinct group IDs (integers or strings) declared in any order, any +set of empty groups anywhere in that order, and any assignment of persons to +the occupied groups, ``join_with_persons`` must give a population where: + +* each person's ``members_entity_id`` points at the group their ID names; +* per-group sums and member counts, in declared order, equal a plain-Python + dict oracle; +* projecting group values gives each person their own group's value. + +Examples, including the audit witness, are in +``test_join_with_persons_identity.py``. +""" + +from __future__ import annotations + +from collections import Counter + +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 tests.core.test_join_with_persons_identity import _join # noqa: E402 + + +@st.composite +def _groups_and_members(draw): + use_text = draw(st.booleans()) + distinct = draw( + st.lists( + st.text(alphabet="abcxyz", min_size=1, max_size=3) + if use_text + else st.integers(min_value=-1000, max_value=1000), + min_size=1, + max_size=10, + unique=True, + ) + ) + group_ids = draw(st.permutations(distinct)) + occupied = draw(st.lists(st.sampled_from(group_ids), min_size=1, unique=True)) + persons_group_ids = draw(st.lists(st.sampled_from(occupied), max_size=25)) + incomes = draw( + st.lists( + st.integers(min_value=-10_000, max_value=10_000), + min_size=len(persons_group_ids), + max_size=len(persons_group_ids), + ) + ) + group_values = draw( + st.lists( + st.integers(min_value=-10_000, max_value=10_000), + min_size=len(group_ids), + max_size=len(group_ids), + ) + ) + return group_ids, persons_group_ids, incomes, group_values + + +@settings(max_examples=300, deadline=None) +@given(_groups_and_members()) +def test_join_matches_dict_oracle(tax_benefit_system, case): + group_ids, persons_group_ids, incomes, group_values = case + household = _join(tax_benefit_system, group_ids, persons_group_ids) + + # Identity: each person sits in the group their ID names. + members = household.members_entity_id + assert [group_ids[i] for i in members] == persons_group_ids + + # Per-group sums and counts, in declared order, match the oracle. + totals = {group_id: 0 for group_id in group_ids} + for group_id, income in zip(persons_group_ids, incomes): + totals[group_id] += income + counts = Counter(persons_group_ids) + assert household.sum(np.array(incomes, dtype=float)).tolist() == [ + float(totals[group_id]) for group_id in group_ids + ] + assert household.nb_persons().tolist() == [ + counts[group_id] for group_id in group_ids + ] + + # Projection gives each person their own group's value. + value_of = dict(zip(group_ids, group_values)) + assert household.project(np.array(group_values)).tolist() == [ + value_of[group_id] for group_id in persons_group_ids + ] From c3a6459cb95218b8dd5632fe4067d2e3fa4686a8 Mon Sep 17 00:00:00 2001 From: Max Ghenis Date: Tue, 6 Oct 2026 16:38:07 -0400 Subject: [PATCH 2/4] Match group IDs by value across dtypes; keep the flat-file fallbacks Addresses the independent review of d647f78d (three should-fix findings). - group_positions compared IDs through np.searchsorted, which promotes mixed dtypes. int64 with float64, or int64 with uint64, meet as float64, where neighbouring large IDs become equal. Groups [2**53+3, 2**53+4] with a float64 person ID 2**53+4 still joined the first group, and int64 IDs [2**53, 2**53+1] with uint64 memberships raised as undeclared. IDs of one kind still use the sorted search (exact within a kind); IDs of different kinds, or object arrays, are now matched one by one through a dict, which compares Python values exactly. - Flat files with a household_id column but no person_household_id column worked on master and raised KeyError here. The membership now comes from the person's membership column or, failing that, the group's ID column repeated on each member's row. A frame with neither raises the same KeyError as master. - Default roles were sized by the number of groups. With the flat file's household count now right, two persons in one household got one role and role-filtered sums raised. They are now one per person (as in #595). Adds regressions for each, and a Hypothesis property: numeric IDs near 0, 2**31, 2**53 and 2**62 match by value for int32/int64/uint64/float64 pairs. All fail on d647f78d; the property shrinks to int64 groups and uint64 memberships near 2**62. The real-dataset mapping A/B (populace-us, enhanced CPS 2024, UK enhanced FRS 1.40.3 and 1.57.4) is still identical. Co-Authored-By: Claude Opus 5.5 --- changelog.d/fix-join-group-identity.fixed.md | 2 +- policyengine_core/simulations/simulation.py | 19 ++++- .../simulations/simulation_builder.py | 80 ++++++++++++++++--- tests/core/test_join_with_persons_identity.py | 67 ++++++++++++++++ ...est_join_with_persons_identity_property.py | 50 ++++++++++++ 5 files changed, 200 insertions(+), 18 deletions(-) diff --git a/changelog.d/fix-join-group-identity.fixed.md b/changelog.d/fix-join-group-identity.fixed.md index 45f7439da..63b8d0994 100644 --- a/changelog.d/fix-join-group-identity.fixed.md +++ b/changelog.d/fix-join-group-identity.fixed.md @@ -1 +1 @@ -Persons joined to groups from a dataset now land in the group their ID names even when a group with no members sorts before theirs; an ID that names no declared group, or a repeated group ID, now raises instead of silently picking a group. A flat-file dataset with a `household_id`-style column no longer declares the previous entity's IDs as its groups. +Persons joined to groups from a dataset now land in the group their ID names even when a group with no members sorts before theirs, and IDs of different dtypes match by value; an ID that names no declared group, or a repeated group ID, now raises instead of silently picking a group. A flat-file dataset with a `household_id`-style column no longer declares the previous entity's IDs as its groups, and its default roles cover every person. diff --git a/policyengine_core/simulations/simulation.py b/policyengine_core/simulations/simulation.py index 8e9ea4f39..fbd2e3a3d 100644 --- a/policyengine_core/simulations/simulation.py +++ b/policyengine_core/simulations/simulation.py @@ -586,9 +586,19 @@ def get_eternity_array(name): person_membership_ids = get_eternity_array(person_membership_id_field) else: # A flat file has one row per person, so it declares each - # group only through its members' IDs. Number the distinct - # IDs 0, 1, ... and give each person their group's number. - membership_values = get_eternity_array(person_membership_id_field) + # group only through its members' rows: their membership + # column or, failing that, the group's ID column repeated on + # each member's row. Number the distinct IDs 0, 1, ... and + # give each person their group's number. + membership_field = next( + ( + field + for field in (person_membership_id_field, entity_id_field) + if any(column.split("__")[0] == field for column in data) + ), + person_membership_id_field, + ) + membership_values = get_eternity_array(membership_field) distinct_ids, person_membership_ids = np.unique( membership_values, return_inverse=True ) @@ -602,7 +612,8 @@ def get_eternity_array(name): elif "role" in data: person_roles = get_eternity_array("role") elif self.default_role is not None: - person_roles = np.full(len(entity_ids), self.default_role) + # One role per person, like the membership IDs. + person_roles = np.full(len(person_membership_ids), self.default_role) else: raise ValueError( f"Missing {person_role_field} column in the dataset. Each group entity must have a person role array defined for ETERNITY." diff --git a/policyengine_core/simulations/simulation_builder.py b/policyengine_core/simulations/simulation_builder.py index b4bba8732..f3590f661 100644 --- a/policyengine_core/simulations/simulation_builder.py +++ b/policyengine_core/simulations/simulation_builder.py @@ -40,36 +40,90 @@ def group_positions( group whose declared ID equals theirs, wherever that group sits in the declared order and whether or not other groups have members. + IDs match only when they are equal as values: IDs of different kinds + (say, 64-bit integers and floats, or signed and unsigned integers) are + compared one by one, because converting them to a common NumPy type can + make different IDs equal. + Raises: ValueError: if the declared IDs repeat, or a person belongs to an ID that was not declared. """ group_ids = np.asarray(group_ids) persons_group_ids = np.asarray(persons_group_ids) + if ( + group_ids.dtype.kind == persons_group_ids.dtype.kind + and group_ids.dtype.kind in "biufUS" + ): + # Within one kind, NumPy compares without losing precision. + return _group_positions_by_sorting(group_ids, persons_group_ids, entity_key) + return _group_positions_by_lookup(group_ids, persons_group_ids, entity_key) + + +def _group_positions_by_sorting( + group_ids: np.ndarray, persons_group_ids: np.ndarray, entity_key: str +) -> np.ndarray: sorter = np.argsort(group_ids, kind="stable") sorted_ids = group_ids[sorter] repeated = sorted_ids[1:][sorted_ids[1:] == sorted_ids[:-1]] if len(repeated) > 0: - raise ValueError( - f"{entity_key} IDs must be unique, but these repeat: " - f"{np.unique(repeated)[:5].tolist()}." + _raise_repeated_group_ids(entity_key, np.unique(repeated).tolist()) + positions = np.searchsorted(sorted_ids, persons_group_ids) + declared = positions < len(sorted_ids) + declared[declared] = sorted_ids[positions[declared]] == persons_group_ids[declared] + if not declared.all(): + _raise_undeclared_group_ids( + entity_key, + int((~declared).sum()), + np.unique(persons_group_ids[~declared]).tolist(), ) + return sorter[positions] + + +def _group_positions_by_lookup( + group_ids: np.ndarray, persons_group_ids: np.ndarray, entity_key: str +) -> np.ndarray: + position_of = {} + repeated = [] try: - positions = np.searchsorted(sorted_ids, persons_group_ids) + for position, group_id in enumerate(group_ids.tolist()): + if group_id in position_of: + repeated.append(group_id) + else: + position_of[group_id] = position + if repeated: + _raise_repeated_group_ids(entity_key, repeated) + positions = np.fromiter( + (position_of.get(group_id, -1) for group_id in persons_group_ids.tolist()), + dtype=np.intp, + count=len(persons_group_ids), + ) except TypeError as error: raise ValueError( f"Person {entity_key} IDs ({persons_group_ids.dtype}) cannot be " - f"compared with the declared {entity_key} IDs ({group_ids.dtype})." + f"matched with the declared {entity_key} IDs ({group_ids.dtype})." ) from error - declared = positions < len(sorted_ids) - declared[declared] = sorted_ids[positions[declared]] == persons_group_ids[declared] - if not declared.all(): - undeclared = np.unique(persons_group_ids[~declared]) - raise ValueError( - f"{int((~declared).sum())} person(s) belong to {entity_key} IDs " - f"that were not declared: {undeclared[:5].tolist()}." + undeclared = positions < 0 + if undeclared.any(): + _raise_undeclared_group_ids( + entity_key, + int(undeclared.sum()), + list(dict.fromkeys(persons_group_ids[undeclared].tolist())), ) - return sorter[positions] + return positions + + +def _raise_repeated_group_ids(entity_key: str, repeated: list) -> None: + raise ValueError( + f"{entity_key} IDs must be unique, but these repeat: {repeated[:5]}." + ) + + +def _raise_undeclared_group_ids(entity_key: str, count: int, undeclared: list) -> None: + raise ValueError( + f"{count} person(s) belong to {entity_key} IDs that were not " + f"declared: {undeclared[:5]}." + ) from datetime import datetime diff --git a/tests/core/test_join_with_persons_identity.py b/tests/core/test_join_with_persons_identity.py index f4f4d7223..80f8f8e6d 100644 --- a/tests/core/test_join_with_persons_identity.py +++ b/tests/core/test_join_with_persons_identity.py @@ -18,8 +18,10 @@ import pytest from policyengine_core.data import Dataset +from policyengine_core.entities import build_entity from policyengine_core.simulations import Simulation, SimulationBuilder from policyengine_core.simulations.simulation_builder import group_positions +from policyengine_core.taxbenefitsystems import TaxBenefitSystem def _join(tax_benefit_system, group_ids, persons_group_ids): @@ -146,3 +148,68 @@ def test_flat_file_with_household_id_column_counts_its_households( household = simulation.populations["household"] assert household.count == 2 assert household.members_entity_id.tolist() == [0, 0, 1] + + +@pytest.mark.parametrize( + "group_ids, persons_group_ids, expected", + [ + # Converting both to float64 rounds 2**53 + 3 up to 2**53 + 4, so the + # person matched the first household. + ( + np.array([2**53 + 3, 2**53 + 4], dtype=np.int64), + np.array([2**53 + 4], dtype=np.float64), + [1], + ), + # int64 and uint64 meet as float64, where 2**53 + 1 becomes 2**53. + ( + np.array([2**53, 2**53 + 1], dtype=np.int64), + np.array([2**53, 2**53 + 1], dtype=np.uint64), + [0, 1], + ), + (np.array([5, 7], dtype=np.int32), np.array([7, 5, 7]), [1, 0, 1]), + ], +) +def test_ids_of_different_dtypes_match_by_value(group_ids, persons_group_ids, expected): + assert group_positions(group_ids, persons_group_ids).tolist() == expected + + +def test_flat_file_with_only_a_household_id_column_reads_it_as_membership( + tax_benefit_system, +): + # With no ``person_household_id`` column, each row's ``household_id`` is + # its person's household. + dataframe = pd.DataFrame({"person_id": [0, 1, 2], "household_id": [7, 7, 9]}) + simulation = Simulation( + tax_benefit_system=tax_benefit_system, + dataset=Dataset.from_dataframe(dataframe, "2024"), + ) + + household = simulation.populations["household"] + assert household.count == 2 + assert household.members_entity_id.tolist() == [0, 0, 1] + + +def test_flat_file_default_roles_cover_every_person(): + # Two persons in one household need two default roles: counting one + # per household gave role-filtered sums a one-entry filter. + person = build_entity("person", "persons", "", is_person=True) + household = build_entity( + "household", "households", "", roles=[{"key": "member", "plural": "members"}] + ) + dataframe = pd.DataFrame( + { + "person_id": [10, 11], + "household_id": [500, 500], + "person_household_id": [500, 500], + } + ) + simulation = Simulation( + tax_benefit_system=TaxBenefitSystem([person, household]), + dataset=Dataset.from_dataframe(dataframe, "2024"), + ) + + population = simulation.populations["household"] + assert len(population.members_role) == 2 + assert population.sum( + np.array([1.0, 2.0]), role=population.entity.MEMBER + ).tolist() == [3.0] diff --git a/tests/core/test_join_with_persons_identity_property.py b/tests/core/test_join_with_persons_identity_property.py index 17c979cc2..6155da2f6 100644 --- a/tests/core/test_join_with_persons_identity_property.py +++ b/tests/core/test_join_with_persons_identity_property.py @@ -9,6 +9,10 @@ dict oracle; * projecting group values gives each person their own group's value. +Numeric IDs also match by value when the declared and person IDs have +different dtypes (int32, int64, uint64, float64), including values near +2**53 and 2**62 where a common float64 would merge neighbours. + Examples, including the audit witness, are in ``test_join_with_persons_identity.py``. """ @@ -23,9 +27,13 @@ # The smoke job installs Core without the dev extra but collects every module. pytest.importorskip("hypothesis") +import hypothesis # noqa: E402 from hypothesis import given, settings # noqa: E402 from hypothesis import strategies as st # noqa: E402 +from policyengine_core.simulations.simulation_builder import ( # noqa: E402 + group_positions, +) from tests.core.test_join_with_persons_identity import _join # noqa: E402 @@ -89,3 +97,45 @@ def test_join_matches_dict_oracle(tax_benefit_system, case): assert household.project(np.array(group_values)).tolist() == [ value_of[group_id] for group_id in persons_group_ids ] + + +DTYPE_PAIRS = [ + (np.int64, np.int64), + (np.int32, np.int64), + (np.int64, np.uint64), + (np.uint64, np.int64), + (np.int64, np.float64), + (np.float64, np.int64), +] + + +def _exact_in(dtype, value): + """Whether ``value`` survives a round trip through ``dtype``.""" + try: + return int(np.array([value], dtype=dtype)[0]) == value + except (OverflowError, ValueError): + return False + + +@settings(max_examples=300, deadline=None) +@given( + dtypes=st.sampled_from(DTYPE_PAIRS), + base=st.sampled_from([0, 2**31 - 16, 2**53 - 16, 2**62]), + offsets=st.lists(st.integers(0, 31), min_size=1, max_size=8, unique=True), + data=st.data(), +) +def test_numeric_ids_match_by_value_across_dtypes(dtypes, base, offsets, data): + group_dtype, person_dtype = dtypes + group_values = [base + offset for offset in offsets] + persons_values = data.draw(st.lists(st.sampled_from(group_values), max_size=12)) + hypothesis.assume( + all(_exact_in(group_dtype, value) for value in group_values) + and all(_exact_in(person_dtype, value) for value in persons_values) + ) + + positions = group_positions( + np.array(group_values, dtype=group_dtype), + np.array(persons_values, dtype=person_dtype), + ) + + assert [group_values[i] for i in positions] == persons_values From 0d2ba5af182cc80cc10cc2c6629542e9788c24ea Mon Sep 17 00:00:00 2001 From: Max Ghenis Date: Tue, 6 Oct 2026 17:36:28 -0400 Subject: [PATCH 3/4] Read flat-file group values by membership; reject missing memberships Addresses the second independent review (of c3a6459c): - should-fix: build_from_dataset took a flat-file column as already per group whenever its length equalled the group count. A flat file has one row per person, so that happens exactly when every group has one person, and then values went in by row order while groups are numbered in ID order. With household_id [9, 7, 8] each person read another's rent ([800, 900, 700] for [900, 700, 800]). Master had this for frames with both ID columns; this PR's household_id fallback newly loaded frames master refused. Group variables in a flat file now always take each group's first member's value. - nit: np.unique merged every NaN membership into one group. A flat file with missing memberships now raises ValueError. - nit: the changelog now says default roles are one per person in every dataset format (every PE-US Microsimulation has no role column); adds an ARRAYS test for that path. The repeated-ID message lists each ID once. - nit: the dtype property also covers uint64 beyond the int64 limit, uint64/float64 pairs and negative IDs. The row-order and NaN regressions fail on c3a6459c and pass here. Co-Authored-By: Claude Opus 5.5 --- changelog.d/fix-join-group-identity.fixed.md | 2 +- policyengine_core/simulations/simulation.py | 15 +++- .../simulations/simulation_builder.py | 2 +- tests/core/test_join_with_persons_identity.py | 84 +++++++++++++++++++ ...est_join_with_persons_identity_property.py | 27 +++++- 5 files changed, 122 insertions(+), 8 deletions(-) diff --git a/changelog.d/fix-join-group-identity.fixed.md b/changelog.d/fix-join-group-identity.fixed.md index 63b8d0994..2693756f1 100644 --- a/changelog.d/fix-join-group-identity.fixed.md +++ b/changelog.d/fix-join-group-identity.fixed.md @@ -1 +1 @@ -Persons joined to groups from a dataset now land in the group their ID names even when a group with no members sorts before theirs, and IDs of different dtypes match by value; an ID that names no declared group, or a repeated group ID, now raises instead of silently picking a group. A flat-file dataset with a `household_id`-style column no longer declares the previous entity's IDs as its groups, and its default roles cover every person. +Persons joined to groups from a dataset now land in the group their ID names even when a group with no members sorts before theirs, and IDs of different dtypes match by value; an ID that names no declared group, or a repeated group ID, now raises instead of silently picking a group. Datasets with no role column now give every person (not every group) the default role, in every dataset format. Flat-file datasets read group values from each group's members rather than by row order, read memberships from a `household_id`-style column when there is no `person_household_id`-style one, reject missing memberships, and no longer declare the previous entity's IDs as their groups. diff --git a/policyengine_core/simulations/simulation.py b/policyengine_core/simulations/simulation.py index fbd2e3a3d..f5d0c9f6d 100644 --- a/policyengine_core/simulations/simulation.py +++ b/policyengine_core/simulations/simulation.py @@ -599,6 +599,13 @@ def get_eternity_array(name): person_membership_id_field, ) membership_values = get_eternity_array(membership_field) + if pd.isna(membership_values).any(): + # np.unique would merge every missing membership into one + # group. + raise ValueError( + f"{int(pd.isna(membership_values).sum())} person(s) " + f"have no {membership_field} in the dataset." + ) distinct_ids, person_membership_ids = np.unique( membership_values, return_inverse=True ) @@ -663,8 +670,12 @@ def get_eternity_array(name): entity = variable_meta.entity population = self.get_population(entity.plural) - # All data should be person level - if len(data[variable]) != len(population.ids): + # A flat file has one row per person, so a group variable's + # column holds each member's copy: take the first member's. + # (Treating a column as already per group whenever its length + # matched the group count put values in row order, not group + # order, when every group had one person.) + if not population.entity.is_person: population: GroupPopulation entity_level_data = population.value_from_first_person( data[variable] diff --git a/policyengine_core/simulations/simulation_builder.py b/policyengine_core/simulations/simulation_builder.py index f3590f661..a68536902 100644 --- a/policyengine_core/simulations/simulation_builder.py +++ b/policyengine_core/simulations/simulation_builder.py @@ -92,7 +92,7 @@ def _group_positions_by_lookup( else: position_of[group_id] = position if repeated: - _raise_repeated_group_ids(entity_key, repeated) + _raise_repeated_group_ids(entity_key, list(dict.fromkeys(repeated))) positions = np.fromiter( (position_of.get(group_id, -1) for group_id in persons_group_ids.tolist()), dtype=np.intp, diff --git a/tests/core/test_join_with_persons_identity.py b/tests/core/test_join_with_persons_identity.py index 80f8f8e6d..d9916afa7 100644 --- a/tests/core/test_join_with_persons_identity.py +++ b/tests/core/test_join_with_persons_identity.py @@ -213,3 +213,87 @@ def test_flat_file_default_roles_cover_every_person(): assert population.sum( np.array([1.0, 2.0]), role=population.entity.MEMBER ).tolist() == [3.0] + + +@pytest.mark.parametrize( + "id_columns", + [ + { + "person_id__2024": [0, 1, 2], + "household_id__2024": [9, 7, 8], + "person_household_id__2024": [9, 7, 8], + }, + {"person_id": [0, 1, 2], "household_id": [9, 7, 8]}, + ], +) +def test_flat_file_group_values_follow_membership_not_row_order( + tax_benefit_system, id_columns +): + # One person per household, IDs not in row order. A household column + # as long as the household count was taken as already per household, + # in row order, while households are numbered in ID order: each person + # read another person's rent. + dataframe = pd.DataFrame({**id_columns, "rent__2024-01": [900.0, 700.0, 800.0]}) + simulation = Simulation( + tax_benefit_system=tax_benefit_system, + dataset=Dataset.from_dataframe(dataframe, "2024"), + ) + + assert simulation.calculate("rent", "2024-01", map_to="person").tolist() == [ + 900.0, + 700.0, + 800.0, + ] + assert simulation.calculate("household_id", "2024", map_to="person").tolist() == [ + 9, + 7, + 8, + ] + + +def test_flat_file_rejects_missing_memberships(tax_benefit_system): + # np.unique counts every NaN as one value, which made one household of + # everyone with no membership. + dataframe = pd.DataFrame( + {"person_id": [0, 1, 2], "person_household_id": [1.0, np.nan, np.nan]} + ) + with pytest.raises(ValueError, match="2 person\\(s\\) have no"): + Simulation( + tax_benefit_system=tax_benefit_system, + dataset=Dataset.from_dataframe(dataframe, "2024"), + ) + + +def test_dataset_default_roles_cover_every_person(tmp_path): + # Every dataset format with no role column got one default role per + # group: three persons in two households had two roles, and a + # role-filtered sum raised IndexError. + person = build_entity("person", "persons", "", is_person=True) + household = build_entity( + "household", "households", "", roles=[{"key": "member", "plural": "members"}] + ) + + class NoRoleColumn(Dataset): + name = "no_role_column" + label = "No role column" + file_path = tmp_path / "no_role_column.h5" + data_format = Dataset.ARRAYS + time_period = "2024" + + NoRoleColumn().save_dataset( + { + "person_id": np.array([1, 2, 3]), + "household_id": np.array([10, 20]), + "person_household_id": np.array([10, 20, 20]), + } + ) + simulation = Simulation( + tax_benefit_system=TaxBenefitSystem([person, household]), + dataset=NoRoleColumn(), + ) + + population = simulation.populations["household"] + assert len(population.members_role) == 3 + assert population.sum( + np.array([1.0, 2.0, 2.0]), role=population.entity.MEMBER + ).tolist() == [1.0, 4.0] diff --git a/tests/core/test_join_with_persons_identity_property.py b/tests/core/test_join_with_persons_identity_property.py index 6155da2f6..5d9b202d7 100644 --- a/tests/core/test_join_with_persons_identity_property.py +++ b/tests/core/test_join_with_persons_identity_property.py @@ -10,8 +10,9 @@ * projecting group values gives each person their own group's value. Numeric IDs also match by value when the declared and person IDs have -different dtypes (int32, int64, uint64, float64), including values near -2**53 and 2**62 where a common float64 would merge neighbours. +different dtypes (int32, int64, uint64, float64), including negative values +and values near 2**31, 2**53, 2**62 and beyond the int64 limit, where a +common float64 would merge neighbours. Examples, including the audit witness, are in ``test_join_with_persons_identity.py``. @@ -104,9 +105,15 @@ def test_join_matches_dict_oracle(tax_benefit_system, case): (np.int32, np.int64), (np.int64, np.uint64), (np.uint64, np.int64), + (np.uint64, np.uint64), (np.int64, np.float64), (np.float64, np.int64), + (np.uint64, np.float64), + (np.float64, np.uint64), ] +#: Where runs of 32 consecutive IDs start: around 0, the int32 limit, the +#: last integers float64 holds exactly, and beyond the int64 limit. +BASES = [-(2**53) - 16, -16, 0, 2**31 - 16, 2**53 - 16, 2**62, 2**63 + 8] def _exact_in(dtype, value): @@ -120,12 +127,24 @@ def _exact_in(dtype, value): @settings(max_examples=300, deadline=None) @given( dtypes=st.sampled_from(DTYPE_PAIRS), - base=st.sampled_from([0, 2**31 - 16, 2**53 - 16, 2**62]), offsets=st.lists(st.integers(0, 31), min_size=1, max_size=8, unique=True), data=st.data(), ) -def test_numeric_ids_match_by_value_across_dtypes(dtypes, base, offsets, data): +def test_numeric_ids_match_by_value_across_dtypes(dtypes, offsets, data): group_dtype, person_dtype = dtypes + base = data.draw( + st.sampled_from( + [ + base + for base in BASES + if all( + _exact_in(dtype, base) and _exact_in(dtype, base + 31) + for dtype in dtypes + if np.dtype(dtype).kind != "f" + ) + ] + ) + ) group_values = [base + offset for offset in offsets] persons_values = data.draw(st.lists(st.sampled_from(group_values), max_size=12)) hypothesis.assume( From cf2eedaee5a05d31ba94a720ae59e61d13b5a812 Mon Sep 17 00:00:00 2001 From: Max Ghenis Date: Sat, 10 Oct 2026 13:35:39 -0400 Subject: [PATCH 4/4] Preserve empty flat-file loading after group identity joins --- changelog.d/fix-join-group-identity.fixed.md | 2 +- policyengine_core/simulations/simulation.py | 4 +- .../simulations/simulation_builder.py | 12 +++ tests/core/test_join_with_persons_identity.py | 68 +++++++++++++- ...est_join_with_persons_identity_property.py | 93 ++++++++++++++++++- 5 files changed, 173 insertions(+), 6 deletions(-) diff --git a/changelog.d/fix-join-group-identity.fixed.md b/changelog.d/fix-join-group-identity.fixed.md index 2693756f1..1bbd009e2 100644 --- a/changelog.d/fix-join-group-identity.fixed.md +++ b/changelog.d/fix-join-group-identity.fixed.md @@ -1 +1 @@ -Persons joined to groups from a dataset now land in the group their ID names even when a group with no members sorts before theirs, and IDs of different dtypes match by value; an ID that names no declared group, or a repeated group ID, now raises instead of silently picking a group. Datasets with no role column now give every person (not every group) the default role, in every dataset format. Flat-file datasets read group values from each group's members rather than by row order, read memberships from a `household_id`-style column when there is no `person_household_id`-style one, reject missing memberships, and no longer declare the previous entity's IDs as their groups. +Correct dataset group identity, membership validation, person-sized default roles, and flat-file membership and value loading while preserving empty datasets. diff --git a/policyengine_core/simulations/simulation.py b/policyengine_core/simulations/simulation.py index f5d0c9f6d..aa05bfd10 100644 --- a/policyengine_core/simulations/simulation.py +++ b/policyengine_core/simulations/simulation.py @@ -675,7 +675,9 @@ def get_eternity_array(name): # (Treating a column as already per group whenever its length # matched the group count put values in row order, not group # order, when every group had one person.) - if not population.entity.is_person: + # With no person rows there are no groups or first members; + # pass the empty column through without member extraction. + if not population.entity.is_person and population.count > 0: population: GroupPopulation entity_level_data = population.value_from_first_person( data[variable] diff --git a/policyengine_core/simulations/simulation_builder.py b/policyengine_core/simulations/simulation_builder.py index a68536902..9ad8daec6 100644 --- a/policyengine_core/simulations/simulation_builder.py +++ b/policyengine_core/simulations/simulation_builder.py @@ -379,6 +379,18 @@ def join_with_persons( persons_group_assignment: ArrayLike, roles: typing.Iterable[str], ) -> None: + """Join persons to the groups named by their membership IDs. + + Memberships and roles are supplied in person order. Each membership + must match a unique declared group ID by value; equal numeric IDs + may have different dtypes. ``members_entity_id`` indexes the declared + group order, including groups with no members. An empty assignment + is supported. Roles are interpreted as role keys or integer indices. + + Raises: + ValueError: if declared group IDs repeat or a membership names + an undeclared group. + """ group_population.members_entity_id = group_positions( group_population.ids, persons_group_assignment, diff --git a/tests/core/test_join_with_persons_identity.py b/tests/core/test_join_with_persons_identity.py index d9916afa7..b6f55729d 100644 --- a/tests/core/test_join_with_persons_identity.py +++ b/tests/core/test_join_with_persons_identity.py @@ -37,6 +37,7 @@ def _join(tax_benefit_system, group_ids, persons_group_ids): def test_person_joins_named_group_when_earlier_group_is_empty(tax_benefit_system): # The audit witness: group "a" is empty and sorts before group "b". + # "b" is at position 1: group sums are [0, 100], and [10, 20][1] = 20. household = _join(tax_benefit_system, ["a", "b"], ["b"]) assert household.members_entity_id.tolist() == [1] @@ -198,9 +199,9 @@ def test_flat_file_default_roles_cover_every_person(): ) dataframe = pd.DataFrame( { - "person_id": [10, 11], - "household_id": [500, 500], - "person_household_id": [500, 500], + "person_id__2024": [10, 11], + "household_id__2024": [500, 500], + "person_household_id__2024": [500, 500], } ) simulation = Simulation( @@ -210,11 +211,70 @@ def test_flat_file_default_roles_cover_every_person(): population = simulation.populations["household"] assert len(population.members_role) == 2 + assert simulation.persons.has_role(population.entity.MEMBER).tolist() == [ + True, + True, + ] + # A single True previously broadcast here; the projected values stay + # the same even though has_role now returns one entry per person. + assert population.project( + np.array([10.0]), role=population.entity.MEMBER + ).tolist() == [10.0, 10.0] + # Both members contribute: 1 + 2 = 3. assert population.sum( np.array([1.0, 2.0]), role=population.entity.MEMBER ).tolist() == [3.0] +@pytest.mark.parametrize("period_suffix", ["", "__2024"]) +def test_empty_flat_file_loads_group_columns(tax_benefit_system, period_suffix): + # Zero person rows declare zero households; both ID and rent inputs + # must stay empty, without asking for a nonexistent first member. + dataframe = pd.DataFrame( + { + f"person_id{period_suffix}": np.array([], dtype=np.int64), + f"household_id{period_suffix}": np.array([], dtype=np.int64), + f"person_household_id{period_suffix}": np.array([], dtype=np.int64), + "rent__2024-01": np.array([], dtype=np.float64), + } + ) + simulation = Simulation( + tax_benefit_system=tax_benefit_system, + dataset=Dataset.from_dataframe(dataframe, "2024"), + ) + + household = simulation.populations["household"] + assert simulation.persons.count == household.count == 0 + assert household.members_entity_id.tolist() == [] + assert simulation.calculate("household_id", "2024").tolist() == [] + assert simulation.calculate("rent", "2024-01").tolist() == [] + assert simulation.calculate("rent", "2024-01", map_to="person").tolist() == [] + + +def test_subsample_fraction_rounding_to_zero_loads_empty_frame(tax_benefit_system): + dataframe = pd.DataFrame( + { + "person_id__2024": [0, 1, 2], + "household_id__2024": [7, 7, 9], + "person_household_id__2024": [7, 7, 9], + "household_weight__2024": [1.0, 1.0, 1.0], + "rent__2024-01": [700.0, 700.0, 900.0], + } + ) + simulation = Simulation( + tax_benefit_system=tax_benefit_system, + dataset=Dataset.from_dataframe(dataframe, "2024"), + ) + + # int(2 households * 0.1) = int(0.2) = 0 selected households. + simulation.subsample(frac=0.1, seed="empty-flat-file", time_period="2024") + + assert simulation.persons.count == simulation.populations["household"].count == 0 + assert simulation.calculate("household_id", "2024").tolist() == [] + assert simulation.calculate("rent", "2024-01").tolist() == [] + assert simulation.calculate("rent", "2024-01", map_to="person").tolist() == [] + + @pytest.mark.parametrize( "id_columns", [ @@ -233,6 +293,8 @@ def test_flat_file_group_values_follow_membership_not_row_order( # as long as the household count was taken as already per household, # in row order, while households are numbered in ID order: each person # read another person's rent. + # Sorted IDs [7, 8, 9] give positions [2, 0, 1] and group rents + # [700, 800, 900], which project back to [900, 700, 800]. dataframe = pd.DataFrame({**id_columns, "rent__2024-01": [900.0, 700.0, 800.0]}) simulation = Simulation( tax_benefit_system=tax_benefit_system, diff --git a/tests/core/test_join_with_persons_identity_property.py b/tests/core/test_join_with_persons_identity_property.py index 5d9b202d7..f6d6b5342 100644 --- a/tests/core/test_join_with_persons_identity_property.py +++ b/tests/core/test_join_with_persons_identity_property.py @@ -14,6 +14,10 @@ and values near 2**31, 2**53, 2**62 and beyond the int64 limit, where a common float64 would merge neighbours. +Flat-file group values match each group's first member, including empty +frames. Default-role masks have one entry per person, and role-filtered +sums and projections match independent membership oracles. + Examples, including the audit witness, are in ``test_join_with_persons_identity.py``. """ @@ -23,18 +27,23 @@ from collections import Counter import numpy as np +import pandas as pd import pytest # The smoke job installs Core without the dev extra but collects every module. pytest.importorskip("hypothesis") import hypothesis # noqa: E402 -from hypothesis import given, settings # noqa: E402 +from hypothesis import example, given, settings # noqa: E402 from hypothesis import strategies as st # noqa: E402 +from policyengine_core.data import Dataset # noqa: E402 +from policyengine_core.entities import build_entity # noqa: E402 +from policyengine_core.simulations import Simulation # noqa: E402 from policyengine_core.simulations.simulation_builder import ( # noqa: E402 group_positions, ) +from policyengine_core.taxbenefitsystems import TaxBenefitSystem # noqa: E402 from tests.core.test_join_with_persons_identity import _join # noqa: E402 @@ -158,3 +167,85 @@ def test_numeric_ids_match_by_value_across_dtypes(dtypes, offsets, data): ) assert [group_values[i] for i in positions] == persons_values + + +@st.composite +def _flat_file_rows(draw): + memberships = draw(st.lists(st.integers(-10, 10), max_size=12)) + rents = draw( + st.lists( + st.integers(0, 10_000), + min_size=len(memberships), + max_size=len(memberships), + ) + ) + return memberships, rents, draw(st.booleans()), draw(st.booleans()) + + +@settings(max_examples=100, deadline=None) +@example(case=([], [], False, True)) +@given(case=_flat_file_rows()) +def test_flat_file_group_values_match_first_member_oracle(tax_benefit_system, case): + memberships, rents, suffixed, explicit_memberships = case + suffix = "__2024" if suffixed else "" + columns = { + f"person_id{suffix}": np.arange(len(memberships)), + f"household_id{suffix}": np.array(memberships, dtype=np.int64), + "rent__2024-01": np.array(rents, dtype=float), + } + if explicit_memberships: + columns[f"person_household_id{suffix}"] = np.array(memberships, dtype=np.int64) + simulation = Simulation( + tax_benefit_system=tax_benefit_system, + dataset=Dataset.from_dataframe(pd.DataFrame(columns), "2024"), + ) + + first_rent = {} + for membership, rent in zip(memberships, rents): + first_rent.setdefault(membership, rent) + group_ids = sorted(first_rent) + assert simulation.persons.count == len(memberships) + assert simulation.populations["household"].count == len(group_ids) + assert simulation.calculate("rent", "2024-01").tolist() == [ + first_rent[group_id] for group_id in group_ids + ] + assert simulation.calculate("rent", "2024-01", map_to="person").tolist() == [ + first_rent[membership] for membership in memberships + ] + + +@settings(max_examples=100, deadline=None) +@example(memberships=[]) +@example(memberships=[7, 7]) +@given(memberships=st.lists(st.integers(-10, 10), max_size=20)) +def test_default_role_masks_cover_persons_and_preserve_role_projection(memberships): + person = build_entity("person", "persons", "", is_person=True) + household = build_entity( + "household", "households", "", roles=[{"key": "member", "plural": "members"}] + ) + dataframe = pd.DataFrame( + { + "person_id__2024": np.arange(len(memberships)), + "person_household_id__2024": np.array(memberships, dtype=np.int64), + } + ) + simulation = Simulation( + tax_benefit_system=TaxBenefitSystem([person, household]), + dataset=Dataset.from_dataframe(dataframe, "2024"), + ) + + population = simulation.populations["household"] + assert population.members_role.shape == (len(memberships),) + assert simulation.persons.has_role(household.MEMBER).tolist() == [ + True for _ in memberships + ] + counts = Counter(memberships) + group_ids = sorted(counts) + assert population.sum( + np.ones(len(memberships)), role=household.MEMBER + ).tolist() == [counts[group_id] for group_id in group_ids] + values = np.arange(len(group_ids)) + 10 + value_of = dict(zip(group_ids, values.tolist())) + assert population.project(values, role=household.MEMBER).tolist() == [ + value_of[membership] for membership in memberships + ]