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..34690433a --- /dev/null +++ b/changelog.d/fix-join-group-identity.fixed.md @@ -0,0 +1 @@ +Correct dataset group identity and membership validation, and load flat-file memberships and group values in group order while preserving empty datasets. diff --git a/policyengine_core/simulations/simulation.py b/policyengine_core/simulations/simulation.py index ddda2ab94..577624eb3 100644 --- a/policyengine_core/simulations/simulation.py +++ b/policyengine_core/simulations/simulation.py @@ -1124,30 +1124,44 @@ 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' 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) + 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 + ) + 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 any(column.split("__")[0] == person_role_field for column in data): @@ -1205,8 +1219,14 @@ 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.) + # 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 e3dee0cc5..9ad8daec6 100644 --- a/policyengine_core/simulations/simulation_builder.py +++ b/policyengine_core/simulations/simulation_builder.py @@ -28,6 +28,104 @@ ) 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. + + 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_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: + 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, list(dict.fromkeys(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"matched with the declared {entity_key} IDs ({group_ids.dtype})." + ) from error + undeclared = positions < 0 + if undeclared.any(): + _raise_undeclared_group_ids( + entity_key, + int(undeclared.sum()), + list(dict.fromkeys(persons_group_ids[undeclared].tolist())), + ) + 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 @@ -281,13 +379,23 @@ 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 - ] + """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, + 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..6ff739a88 --- /dev/null +++ b/tests/core/test_join_with_persons_identity.py @@ -0,0 +1,413 @@ +"""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.entities import build_entity +from policyengine_core.periods import period +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): + 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". + # "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] + 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] + + +@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__2024": [10, 11], + "household_id__2024": [500, 500], + "person_household_id__2024": [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 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_flat_file_explicit_roles_override_defaults(period_suffix): + person = build_entity("person", "persons", "", is_person=True) + household = build_entity( + "household", "households", "", roles=[{"key": "member"}, {"key": "guest"}] + ) + dataframe = pd.DataFrame( + { + "person_id__2024": [10, 11], + "person_household_id__2024": [500, 500], + f"person_household_role{period_suffix}": ["member", "guest"], + } + ) + simulation = Simulation( + tax_benefit_system=TaxBenefitSystem([person, household]), + dataset=Dataset.from_dataframe(dataframe, "2024"), + ) + + population = simulation.populations["household"] + assert simulation.persons.has_role(household.MEMBER).tolist() == [True, False] + # Only the first person's income belongs to MEMBER: 1 + 0 = 1. + assert population.sum(np.array([1.0, 2.0]), role=household.MEMBER).tolist() == [1.0] + assert population.project(np.array([10.0]), role=household.MEMBER).tolist() == [ + 10.0, + 0.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() == [] + + +@pytest.mark.parametrize("quantize_weights", [False, True]) +def test_subsample_fraction_rounding_to_zero_rejects_nonzero_weight( + tax_benefit_system, quantize_weights +): + 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], + } + ) + dataset = Dataset.from_dataframe(dataframe, "2024") + simulation = Simulation(tax_benefit_system=tax_benefit_system, dataset=dataset) + household = simulation.populations["household"] + persons = simulation.persons + + # Native household weights are [1, 1], total 1 + 1 = 2, not the + # repeated person-row total 3. int(2 * 0.1) = 0 selects no households; + # retained weight 0 cannot preserve 2, so reject before rebuilding. + with pytest.raises(ValueError, match="household_weight__2024 total.*zero weight"): + simulation.subsample( + frac=0.1, + seed="empty-flat-file", + time_period="2024", + quantize_weights=quantize_weights, + ) + + assert simulation.dataset is dataset + assert simulation.persons is persons + assert simulation.populations["household"] is household + assert simulation.persons.count == 3 + assert household.count == 2 + assert household.members_entity_id.tolist() == [0, 0, 1] + assert simulation.calculate("household_id", "2024").tolist() == [7, 9] + assert simulation.calculate("household_weight", "2024").tolist() == [1.0, 1.0] + assert simulation.calculate("rent", "2024-01").tolist() == [700.0, 900.0] + assert simulation._get_recorded_input_array( + "household_weight", period("2024") + ).tolist() == [1.0, 1.0] + assert simulation._get_recorded_input_array("rent", period("2024-01")).tolist() == [ + 700.0, + 900.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. + # 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, + 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 new file mode 100644 index 000000000..04c03c491 --- /dev/null +++ b/tests/core/test_join_with_persons_identity_property.py @@ -0,0 +1,359 @@ +"""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. + +Numeric IDs also match by value when the declared and person IDs have +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. + +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. + +Explicit entity roles override defaults with or without a period suffix. +Zero-household selections reject positive source weight totals without +replacing the dataset, populations, memberships, or recorded input values. + +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 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 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.periods import period # 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 + + +@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 + ] + + +DTYPE_PAIRS = [ + (np.int64, np.int64), + (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): + """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), + 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, 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( + 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 + + +@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 + ] + + +@settings(max_examples=100, deadline=None) +@example(rows=[(7, True), (7, False)], suffixed=True) +@example(rows=[], suffixed=True) +@given( + rows=st.lists(st.tuples(st.integers(-10, 10), st.booleans()), max_size=20), + suffixed=st.booleans(), +) +def test_explicit_role_masks_match_person_rows(rows, suffixed): + person = build_entity("person", "persons", "", is_person=True) + household = build_entity( + "household", "households", "", roles=[{"key": "member"}, {"key": "guest"}] + ) + memberships = [membership for membership, _ in rows] + role_suffix = "__2024" if suffixed else "" + dataframe = pd.DataFrame( + { + "person_id__2024": np.arange(len(rows)), + "person_household_id__2024": np.array(memberships, dtype=np.int64), + f"person_household_role{role_suffix}": np.array( + ["member" if is_member else "guest" for _, is_member in rows], + dtype=str, + ), + } + ) + simulation = Simulation( + tax_benefit_system=TaxBenefitSystem([person, household]), + dataset=Dataset.from_dataframe(dataframe, "2024"), + ) + + population = simulation.populations["household"] + assert simulation.persons.has_role(household.MEMBER).tolist() == [ + is_member for _, is_member in rows + ] + counts = Counter(membership for membership, is_member in rows if is_member) + group_ids = sorted(set(memberships)) + assert population.sum(np.ones(len(rows)), role=household.MEMBER).tolist() == [ + counts[group_id] for group_id in group_ids + ] + value_of = {group_id: index + 10 for index, group_id in enumerate(group_ids)} + assert population.project( + np.array(list(value_of.values())), role=household.MEMBER + ).tolist() == [ + value_of[membership] if is_member else 0 for membership, is_member in rows + ] + + +@settings(max_examples=100, deadline=None) +@example(memberships=[7, 7, 9], weight=1, quantize_weights=False) +@example(memberships=[7, 7, 9], weight=1, quantize_weights=True) +@given( + memberships=st.lists(st.integers(-10, 10), min_size=1, max_size=20), + weight=st.integers(1, 100), + quantize_weights=st.booleans(), +) +def test_zero_selection_rejects_positive_weight_without_rebuilding( + tax_benefit_system, memberships, weight, quantize_weights +): + dataframe = pd.DataFrame( + { + "person_id__2024": np.arange(len(memberships)), + "household_id__2024": memberships, + "person_household_id__2024": memberships, + "household_weight__2024": np.full(len(memberships), weight, dtype=float), + "rent__2024-01": np.arange(len(memberships), dtype=float), + } + ) + dataset = Dataset.from_dataframe(dataframe, "2024") + simulation = Simulation(tax_benefit_system=tax_benefit_system, dataset=dataset) + household = simulation.populations["household"] + persons = simulation.persons + before_members = household.members_entity_id.copy() + before_rents = simulation.get_holder("rent").get_array(period("2024-01")).copy() + before_weights = ( + simulation.get_holder("household_weight").get_array(period("2024")).copy() + ) + group_count = len(set(memberships)) + + # Source total = group_count * weight > 0. For this fraction, + # int(group_count * (0.5 / group_count)) = int(0.5) = 0; + # the retained total is 0 and cannot conserve the positive source total. + with pytest.raises(ValueError, match="household_weight__2024 total.*zero weight"): + simulation.subsample( + frac=0.5 / group_count, + seed="empty-flat-file-property", + time_period="2024", + quantize_weights=quantize_weights, + ) + + assert simulation.dataset is dataset + assert simulation.persons is persons + assert simulation.populations["household"] is household + assert simulation.persons.count == len(memberships) + assert household.count == group_count + np.testing.assert_array_equal(household.members_entity_id, before_members) + np.testing.assert_array_equal( + simulation._get_recorded_input_array("rent", period("2024-01")), before_rents + ) + np.testing.assert_array_equal( + simulation._get_recorded_input_array("household_weight", period("2024")), + before_weights, + )