diff --git a/changelog.d/fix-default-role-person-length.fixed.md b/changelog.d/fix-default-role-person-length.fixed.md new file mode 100644 index 000000000..5a7c7973b --- /dev/null +++ b/changelog.d/fix-default-role-person-length.fixed.md @@ -0,0 +1 @@ +Reject implicit dataset defaults that exceed their recognized role's capacity per group and role arrays that are not one dimensional or do not match the membership count. diff --git a/policyengine_core/simulations/simulation.py b/policyengine_core/simulations/simulation.py index ddda2ab94..3d3950b8f 100644 --- a/policyengine_core/simulations/simulation.py +++ b/policyengine_core/simulations/simulation.py @@ -1078,7 +1078,19 @@ def build_from_populations(self, populations: Dict[str, Population]) -> None: self.create_shortcuts() def build_from_dataset(self) -> None: - """Build a simulation from a dataset.""" + """Build a simulation from a dataset. + + Entity-specific role columns take precedence over a generic ``role`` + column, including period-suffixed entity role columns in flat files. + Without either, each person receives the literal ``default_role`` + key. Only flattened-role keys match: unmatched and compound keys remain + unrecognized rather than selecting a first role or subrole. + + A recognized implicit default with a finite role maximum must fit every + group's membership count. Otherwise an explicit role column is required. + Explicit dataset role columns do not validate role capacity. + ``default_role=None`` requires a role column. + """ self.build_from_populations(self.tax_benefit_system.instantiate_entities()) from policyengine_core.simulations.simulation_builder import ( SimulationBuilder, @@ -1155,6 +1167,26 @@ def get_eternity_array(name): elif "role" in data: person_roles = get_eternity_array("role") elif self.default_role is not None: + default_role = next( + ( + role + for role in group_entity.flattened_roles + if role.key == self.default_role + ), + None, + ) + if default_role is not None and default_role.max is not None: + membership_counts = np.unique( + person_membership_ids, return_counts=True + )[1] + if np.any(membership_counts > default_role.max): + raise ValueError( + f"Cannot assign default role {self.default_role!r} " + f"to every person in {group_entity.key}: at most " + f"{default_role.max} member(s) per group may hold it; " + f"provide an explicit {person_role_field} column." + ) + # Preserve master's one literal default key per declared person. person_roles = np.full(self.persons.count, self.default_role) else: raise ValueError( diff --git a/policyengine_core/simulations/simulation_builder.py b/policyengine_core/simulations/simulation_builder.py index e3dee0cc5..7480fa9c0 100644 --- a/policyengine_core/simulations/simulation_builder.py +++ b/policyengine_core/simulations/simulation_builder.py @@ -281,6 +281,13 @@ def join_with_persons( persons_group_assignment: ArrayLike, roles: typing.Iterable[str], ) -> None: + """Join persons to groups with one role per membership. + + ``roles`` must be one dimensional and match the membership array's length. + Integers index the entity's flattened roles; strings match flattened-role + keys literally, with unmatched keys represented as unrecognized roles. + This method does not validate role capacity. + """ # 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 @@ -291,6 +298,14 @@ def join_with_persons( flattened_roles = group_population.entity.flattened_roles roles_array = np.array(roles) + if roles_array.ndim != 1 or len(roles_array) != len(persons_group_assignment): + # ``members_role`` is indexed like ``members_entity_id``: one entry + # per membership. Other shapes can make queries fail or broadcast incorrectly. + raise ValueError( + f"Got {roles_array.size} {group_population.entity.key} role(s) " + f"for {len(persons_group_assignment)} person(s); give one role " + "per person." + ) if np.issubdtype(roles_array.dtype, np.integer): group_population.members_role = np.array(flattened_roles)[roles_array] else: diff --git a/tests/core/test_dataset_default_role.py b/tests/core/test_dataset_default_role.py new file mode 100644 index 000000000..c943628c7 --- /dev/null +++ b/tests/core/test_dataset_default_role.py @@ -0,0 +1,456 @@ +"""A dataset without role columns gives every person the default role. + +``Simulation.build_from_dataset`` falls back to ``Simulation.default_role`` +when a dataset has no ``person__role`` or ``role`` column. It used to +size that fallback by the number of groups, so ``members_role`` had one entry +per group instead of one per person, and role queries (``nb_persons``, +``sum``, ``project``, ``has_role`` with a role) raised or broadcast a +single group's entry over everyone. Literal default keys keep their existing +meaning: unmatched and compound keys leave members holding no recognized role. +Recognized defaults must fit the role's per-group capacity. + +``test_dataset_default_role_property.py`` checks the same invariants for +random datasets, group entities and default roles. +""" + +from __future__ import annotations + +import numpy as np +import pandas as pd +import pytest + +from policyengine_core.country_template import entities +from policyengine_core.data import Dataset +from policyengine_core.entities import build_entity +from policyengine_core.periods import ETERNITY +from policyengine_core.simulations import Simulation, SimulationBuilder +from policyengine_core.taxbenefitsystems import TaxBenefitSystem +from policyengine_core.tools import simulation_dumper + +Household = entities.Household + + +def in_memory_dataset(data, data_format=Dataset.ARRAYS, time_period="2024"): + """A dataset whose ``load`` returns ``data``, like ``Dataset.from_dataframe``.""" + return type( + "InMemoryDataset", + (Dataset,), + { + "name": "in_memory", + "label": "In-memory dataset", + "data_format": data_format, + "file_path": "in_memory", + "time_period": time_period, + "load": lambda self: data, + }, + )() + + +def simulate(tax_benefit_system, dataset, default_role=Simulation.default_role): + simulation_class = type( + "DefaultRoleSimulation", (Simulation,), {"default_role": default_role} + ) + return simulation_class(tax_benefit_system=tax_benefit_system, dataset=dataset) + + +@pytest.mark.parametrize("default_role", ["member", "parent"]) +@pytest.mark.parametrize( + "query", ["value_from_person", "unique_role_projector", "partner_projection"] +) +def test_unmatched_default_preserves_unique_role_queries( + tax_benefit_system, default_role, query +): + # Review witness: both people belong to 10; household 20 is empty. + simulation = simulate( + tax_benefit_system, + in_memory_dataset( + { + "person_id": np.array([1, 2]), + "household_id": np.array([10, 20]), + "person_household_id": np.array([10, 10]), + } + ), + default_role=default_role, + ) + household = simulation.populations["household"] + values = np.array([100, 200]) + if query == "value_from_person": + result = household.value_from_person(values, Household.FIRST_PARENT) + elif query == "unique_role_projector": + result = household.first_parent.transform(values) + else: + result = simulation.persons.value_from_partner( + values, simulation.persons.household, Household.PARENT + ) + # Neither literal key is a flattened role: no parent contributes a value. + assert result.tolist() == [0, 0] + + +@pytest.mark.parametrize("default_role", ["first_parent", "second_parent"]) +def test_unique_default_requires_explicit_roles_when_capacity_exceeded( + tax_benefit_system, default_role +): + with pytest.raises( + ValueError, match="default role.*at most 1.*person_household_role" + ): + simulate( + tax_benefit_system, + in_memory_dataset( + { + "person_id": np.array([1, 2]), + "household_id": np.array([10, 20]), + "person_household_id": np.array([10, 10]), + } + ), + default_role=default_role, + ) + + +def test_unique_default_with_one_person_per_group_supports_projection( + tax_benefit_system, +): + simulation = simulate( + tax_benefit_system, + in_memory_dataset( + { + "person_id": np.array([1, 2]), + "household_id": np.array([10, 20, 30]), + "person_household_id": np.array([20, 10]), + } + ), + default_role="first_parent", + ) + household = simulation.household + values = np.array([100, 200]) + # ID 10 has value 200; ID 20 has value 100; ID 30 has no parent (0). + assert household.value_from_person(values, Household.FIRST_PARENT).tolist() == [ + 200, + 100, + 0, + ] + assert household.first_parent.transform(values).tolist() == [200, 100, 0] + assert simulation.persons.value_from_partner( + values, simulation.persons.household, Household.PARENT + ).tolist() == [0, 0] + + +@pytest.mark.parametrize("capacity, persons", [(0, 0), (0, 1), (2, 2), (2, 3)]) +def test_default_respects_noncompound_role_capacity(capacity, persons): + unit = build_entity( + key="unit", + plural="units", + label="Unit", + roles=[{"key": "member", "max": capacity}], + ) + dataset = in_memory_dataset( + { + "person_id": np.arange(persons), + "unit_id": np.array([10, 20]), + "person_unit_id": np.full(persons, 10, dtype=int), + } + ) + system = TaxBenefitSystem([entities.Person, unit]) + if persons > capacity: + with pytest.raises(ValueError, match=f"default role.*at most {capacity}"): + simulate(system, dataset) + else: + simulation = simulate(system, dataset) + assert list(simulation.unit.members_role) == [unit.MEMBER] * persons + + +@pytest.mark.parametrize( + "default_role, encoded", [("child", "child"), ("member", "unknown")] +) +def test_dataset_default_roles_survive_dump_restore( + tax_benefit_system, tmp_path, default_role, encoded +): + simulation = simulate( + tax_benefit_system, + in_memory_dataset( + { + "person_id": np.array([1, 2, 3]), + "household_id": np.array([10, 20]), + "person_household_id": np.array([10, 10, 20]), + } + ), + default_role=default_role, + ) + directory = tmp_path / "dump" + simulation_dumper.dump_simulation(simulation, str(directory)) + stored_roles = np.load( + directory / "__entities__" / "household" / "members_role.npy" + ) + assert stored_roles.tolist() == [encoded] * 3 + restored = simulation_dumper.restore_simulation(str(directory), tax_benefit_system) + np.testing.assert_array_equal( + restored.household.members_role, simulation.household.members_role + ) + np.testing.assert_array_equal( + restored.household.nb_persons(role=Household.CHILD), + simulation.household.nb_persons(role=Household.CHILD), + ) + + +def test_one_household_of_three_gives_three_roles(tax_benefit_system): + # The audit witness: three persons in one household, no role column. + simulation = simulate( + tax_benefit_system, + in_memory_dataset( + { + "person_id": np.array([1, 2, 3]), + "household_id": np.array([10]), + "person_household_id": np.array([10, 10, 10]), + } + ), + ) + household = simulation.populations["household"] + + assert len(household.members_role) == 3 + # As on master, "member" names no flattened household role. + assert list(household.members_role) == [0] * 3 + assert household.nb_persons(role=Household.PARENT).tolist() == [0] + assert household.nb_persons(role=Household.CHILD).tolist() == [0] + assert household.sum(np.array([1.0, 2.0, 4.0]), role=Household.PARENT).tolist() == [ + # No recognized parents: the filtered sum is 0 + 0 + 0 = 0. + 0.0 + ] + assert simulation.persons.has_role(Household.FIRST_PARENT).tolist() == [False] * 3 + + +def test_saved_dataset_gives_one_role_per_person(tax_benefit_system, tmp_path): + # The same failure through an HDF5 file, with two households of + # different sizes so no shape broadcasts by accident. + class TwoHouseholds(Dataset): + name = "two_households" + label = "Two households" + file_path = tmp_path / "two_households.h5" + data_format = Dataset.TIME_PERIOD_ARRAYS + + TwoHouseholds().save_dataset( + { + "person_id": {ETERNITY: np.array([1, 2, 3, 4, 5])}, + "household_id": {ETERNITY: np.array([10, 20])}, + "person_household_id": {ETERNITY: np.array([20, 10, 20, 20, 10])}, + } + ) + simulation = simulate(tax_benefit_system, TwoHouseholds(), default_role="child") + household = simulation.populations["household"] + + assert list(household.members_role) == [Household.CHILD] * 5 + # ID 10 has persons 1 and 4 (2); ID 20 has persons 0, 2 and 3 (3). + assert household.nb_persons(role=Household.CHILD).tolist() == [2, 3] + assert household.nb_persons(role=Household.PARENT).tolist() == [0, 0] + assert household.project(np.array([7.0, 9.0]), role=Household.CHILD).tolist() == [ + 9.0, + 7.0, + 9.0, + 9.0, + 7.0, + ] + # max(1, 4) = 4 for ID 10; max(0, 2, 3) = 3 for ID 20. + assert household.max(np.arange(5.0), role=Household.CHILD).tolist() == [4.0, 3.0] + + +def test_flat_file_gives_one_role_per_person(tax_benefit_system): + dataframe = pd.DataFrame( + { + "person_id": [10, 11, 12, 13], + "person_household_id": [700, 500, 700, 700], + } + ) + simulation = simulate( + tax_benefit_system, + Dataset.from_dataframe(dataframe, "2024"), + default_role="child", + ) + household = simulation.populations["household"] + + assert len(household.members_role) == 4 + assert household.nb_persons(role=Household.CHILD).tolist() == [1, 3] + + +@pytest.mark.parametrize( + "default_role, expected_role", + [ + ("child", Household.CHILD), + ("second_parent", Household.SECOND_PARENT), + ("parent", 0), + ("member", 0), + ], +) +def test_default_role_keeps_literal_join_semantics( + tax_benefit_system, default_role, expected_role +): + simulation = simulate( + tax_benefit_system, + in_memory_dataset( + { + "person_id": np.array([1, 2]), + "household_id": np.array([10, 20]), + "person_household_id": np.array([20, 10]), + } + ), + default_role=default_role, + ) + assert list(simulation.household.members_role) == [expected_role] * 2 + + +def test_default_role_keeps_literal_semantics_for_each_entity(): + family = build_entity( + key="family", + plural="families", + label="Family", + roles=[ + {"key": "child", "plural": "children"}, + {"key": "parent", "plural": "parents", "subroles": ["mother", "father"]}, + ], + ) + unit = build_entity( + key="unit", + plural="units", + label="Unit", + roles=[{"key": "member", "plural": "members"}], + ) + + simulation = simulate( + TaxBenefitSystem([entities.Person, family, unit]), + in_memory_dataset( + { + "person_id": np.array([1, 2]), + "family_id": np.array([10]), + "person_family_id": np.array([10, 10]), + "unit_id": np.array([20]), + "person_unit_id": np.array([20, 20]), + } + ), + default_role="member", + ) + assert list(simulation.family.members_role) == [0, 0] + assert list(simulation.unit.members_role) == [unit.MEMBER] * 2 + + +def test_no_default_role_still_requires_a_role_column(tax_benefit_system): + with pytest.raises(ValueError, match="person_household_role"): + simulate( + tax_benefit_system, + in_memory_dataset( + { + "person_id": np.array([1, 2]), + "household_id": np.array([10]), + "person_household_id": np.array([10, 10]), + } + ), + default_role=None, + ) + + +@pytest.mark.parametrize("role_column", ["person_household_role", "role"]) +def test_role_column_overrides_the_default(tax_benefit_system, role_column): + simulation = simulate( + tax_benefit_system, + in_memory_dataset( + { + "person_id": np.array([1, 2, 3]), + "household_id": np.array([10]), + "person_household_id": np.array([10, 10, 10]), + # Indices into the flattened roles: first_parent, child, child. + role_column: np.array([0, 2, 2]), + } + ), + # Three people exceed this default's capacity; explicit roles override it. + default_role="first_parent", + ) + household = simulation.populations["household"] + + assert household.nb_persons(role=Household.PARENT).tolist() == [1] + assert household.nb_persons(role=Household.CHILD).tolist() == [2] + + +@pytest.mark.parametrize("roles", [["first_parent"], [["first_parent"]] * 3, "child"]) +def test_join_rejects_invalid_role_shapes(tax_benefit_system, roles): + builder = SimulationBuilder() + builder.create_entities(tax_benefit_system) + builder.declare_person_entity("person", [1, 2, 3]) + household = builder.declare_entity("household", [10]) + + with pytest.raises(ValueError, match="household role\\(s\\) for 3 person"): + builder.join_with_persons(household, [10, 10, 10], roles) + + +def test_three_person_child_default_matches_current_master(tax_benefit_system): + simulation = simulate( + tax_benefit_system, + in_memory_dataset( + { + "person_id": np.array([1, 2, 3]), + "household_id": np.array([10]), + "person_household_id": np.array([10, 10, 10]), + } + ), + default_role="child", + ) + assert list(simulation.household.members_role) == [Household.CHILD] * 3 + # Three memberships in ID 10: 1 + 1 + 1 = 3 children. + assert simulation.household.nb_persons(role=Household.CHILD).tolist() == [3] + + +@pytest.mark.parametrize("role_column", [None, "person_unit_role", "role"]) +def test_bounded_default_rejects_previously_queryable_assignment(role_column): + unit = build_entity( + key="unit", + plural="units", + label="Unit", + roles=[{"key": "member", "max": 2}], + ) + columns = { + "person_id": np.array([1, 2, 3]), + "unit_id": np.array([10]), + "person_unit_id": np.array([10, 10, 10]), + } + if role_column is not None: + columns[role_column] = np.array([0, 0, 0]) + system = TaxBenefitSystem([entities.Person, unit]) + dataset = in_memory_dataset(columns) + if role_column is None: + # Uniform assignment would produce 3 holders, exceeding max=2. + with pytest.raises( + ValueError, match="default role.*at most 2.*person_unit_role" + ): + simulate(system, dataset) + return + + # Capacity enforcement remains limited to synthesized defaults. + simulation = simulate(system, dataset) + assert list(simulation.unit.members_role) == [unit.MEMBER] * 3 + assert simulation.unit.nb_persons(role=unit.MEMBER).tolist() == [3] + # All three explicit members contribute: 1 + 2 + 3 = 6. + assert simulation.unit.sum(np.array([1, 2, 3]), role=unit.MEMBER).tolist() == [6] + + +@pytest.mark.parametrize("generic_role", [False, True]) +def test_suffixed_entity_role_column_overrides_generic_and_default( + tax_benefit_system, generic_role +): + columns = { + "person_id": [1, 2, 3], + "person_household_id": [0, 0, 0], + "person_household_role__2024": [0, 2, 2], + } + if generic_role: + columns["role"] = [2, 2, 2] + simulation = simulate( + tax_benefit_system, + Dataset.from_dataframe(pd.DataFrame(columns), "2024"), + # Ignoring the valid explicit roles would exceed this default's max=1. + default_role="first_parent", + ) + household = simulation.household + assert list(household.members_role) == [ + Household.FIRST_PARENT, + Household.CHILD, + Household.CHILD, + ] + assert household.nb_persons(role=Household.PARENT).tolist() == [1] + assert household.nb_persons(role=Household.CHILD).tolist() == [2] + # The children contribute 2 + 3 = 5. + assert household.sum(np.array([1, 2, 3]), role=Household.CHILD).tolist() == [5] diff --git a/tests/core/test_dataset_default_role_property.py b/tests/core/test_dataset_default_role_property.py new file mode 100644 index 000000000..36d80789c --- /dev/null +++ b/tests/core/test_dataset_default_role_property.py @@ -0,0 +1,388 @@ +"""Dataset defaults preserve literal role matching and person-level cardinality. + +The role table below specifies expected assignments independently of the loader: +only a literal flattened-role key matches; unknown and compound keys remain 0. +For accepted defaults, role counts, aggregates, projections and unique lookups +are computed from the original dataset memberships. A capacity-limited default +must be rejected if assigning it uniformly would exceed its per-group maximum. +""" + +from __future__ import annotations + +from collections import Counter +from typing import NamedTuple + +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") + +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 ETERNITY # noqa: E402 +from policyengine_core.projectors import UniqueRoleToEntityProjector # noqa: E402 +from policyengine_core.simulations import SimulationBuilder # noqa: E402 +from policyengine_core.taxbenefitsystems import TaxBenefitSystem # noqa: E402 +from tests.core.test_dataset_default_role import ( # noqa: E402 + in_memory_dataset, + simulate, +) + +PERSON = build_entity(key="person", plural="persons", label="Person", is_person=True) +HOUSEHOLD = build_entity( + key="household", + plural="households", + label="Household", + roles=[ + { + "key": "parent", + "plural": "parents", + "subroles": ["first_parent", "second_parent"], + }, + {"key": "child", "plural": "children"}, + ], +) +FAMILY = build_entity( + key="family", + plural="families", + label="Family", + roles=[ + {"key": "child", "plural": "children"}, + {"key": "parent", "plural": "parents", "subroles": ["mother", "father"]}, + ], +) +UNIT = build_entity( + key="unit", + plural="units", + label="Unit", + roles=[{"key": "member", "plural": "members"}], +) +COMMITTEE = build_entity( + key="committee", + plural="committees", + label="Committee", + roles=[{"key": "delegate", "plural": "delegates", "max": 2}], +) +GROUP_ENTITIES = [HOUSEHOLD, FAMILY, UNIT, COMMITTEE] +SYSTEM = TaxBenefitSystem([PERSON, *GROUP_ENTITIES]) + +# Deliberately explicit: this oracle does not reproduce a resolution algorithm. +EXPECTED_DEFAULT_ROLES = { + "household": { + "child": HOUSEHOLD.CHILD, + "first_parent": HOUSEHOLD.FIRST_PARENT, + "second_parent": HOUSEHOLD.SECOND_PARENT, + }, + "family": {"child": FAMILY.CHILD, "mother": FAMILY.MOTHER, "father": FAMILY.FATHER}, + "unit": {"member": UNIT.MEMBER}, + "committee": {"delegate": COMMITTEE.DELEGATE}, +} +DEFAULT_ROLES = [ + "member", + "parent", + "child", + "first_parent", + "second_parent", + "mother", + "father", + "delegate", + "nobody", +] + + +class DatasetCase(NamedTuple): + dataset: Dataset + memberships: dict[str, np.ndarray] + values: np.ndarray + + +def _case(columns, data_format, values): + memberships = { + entity.key: columns[f"person_{entity.key}_id"].copy() + for entity in GROUP_ENTITIES + } + if data_format == Dataset.FLAT_FILE: + dataset = Dataset.from_dataframe(pd.DataFrame(columns), "2024") + elif data_format == Dataset.TIME_PERIOD_ARRAYS: + dataset = in_memory_dataset( + {name: {ETERNITY: array} for name, array in columns.items()}, data_format + ) + else: + dataset = in_memory_dataset(columns, data_format) + return DatasetCase(dataset, memberships, np.array(values, dtype=float)) + + +def _uniform_case(persons, group_ids, memberships): + columns = {"person_id": np.arange(1, persons + 1)} + for entity in GROUP_ENTITIES: + columns[f"{entity.key}_id"] = np.array(group_ids, dtype=int) + columns[f"person_{entity.key}_id"] = np.array(memberships, dtype=int) + return _case(columns, Dataset.ARRAYS, np.arange(1, persons + 1) * 100) + + +@st.composite +def _datasets(draw): + persons = draw(st.integers(min_value=0, max_value=12)) + data_format = draw( + st.sampled_from([Dataset.ARRAYS, Dataset.TIME_PERIOD_ARRAYS, Dataset.FLAT_FILE]) + ) + columns = {"person_id": np.arange(100, 100 + persons)} + for entity in GROUP_ENTITIES: + if data_format == Dataset.FLAT_FILE: + # The loader synthesizes group IDs 0..n-1 for a flat file. Choose + # consecutive occupied IDs to isolate roles from ID renumbering. + group_ids = list(range(max(persons, 1))) + else: + group_ids = draw( + st.lists( + st.integers(min_value=-50, max_value=50), + min_size=1 if persons else 0, + max_size=6, + unique=True, + ) + ) + if persons: + occupied = draw( + st.integers(min_value=1, max_value=min(persons, len(group_ids))) + ) + # Keep empty groups after occupied ones in sorted ID order. The + # existing join's treatment of other empty-group positions is + # separate from default-role assignment. + occupied_ids = sorted(group_ids)[:occupied] + memberships = occupied_ids + draw( + st.lists( + st.sampled_from(occupied_ids), + min_size=persons - occupied, + max_size=persons - occupied, + ) + ) + memberships = draw(st.permutations(memberships)) + else: + memberships = [] + if data_format != Dataset.FLAT_FILE: + columns[f"{entity.key}_id"] = np.array(group_ids, dtype=int) + columns[f"person_{entity.key}_id"] = np.array(memberships, dtype=int) + values = draw( + st.lists( + st.integers(min_value=-1000, max_value=1000), + min_size=persons, + max_size=persons, + ) + ) + return _case(columns, data_format, values) + + +@settings(max_examples=200, deadline=None) +@given(_datasets(), st.sampled_from(DEFAULT_ROLES)) +@example(_uniform_case(2, [10, 20], [10, 10]), "member") +@example(_uniform_case(3, [10, 20], [10, 10, 10]), "first_parent") +@example(_uniform_case(3, [10, 20], [10, 10, 10]), "delegate") +@example(_uniform_case(2, [10, 20], [10, 10]), "delegate") +@example(_uniform_case(3, [20, 10, 30], [20, 10, 30]), "first_parent") +def test_default_roles_preserve_literal_matching_and_capacity(case, default_role): + dataset, memberships, values = case + persons = len(values) + violations = [] + for entity in GROUP_ENTITIES: + role = EXPECTED_DEFAULT_ROLES[entity.key].get(default_role) + if role is not None and role.max is not None: + if any( + count > role.max for count in Counter(memberships[entity.key]).values() + ): + violations.append(entity) + if violations: + with pytest.raises(ValueError, match="default role.*at most") as error: + simulate(SYSTEM, dataset, default_role=default_role) + assert f"person_{violations[0].key}_role" in str(error.value) + return + + simulation = simulate(SYSTEM, dataset, default_role=default_role) + + for entity in GROUP_ENTITIES: + group = simulation.populations[entity.key] + role = EXPECTED_DEFAULT_ROLES[entity.key].get(default_role) + raw_memberships = memberships[entity.key] + + assert group.members_role.shape == (persons,) + assert len(group.members_entity_id) == persons + if role is None: + assert (group.members_role == 0).all() + else: + assert all(member_role is role for member_role in group.members_role) + if role.max is not None: + assert all( + count <= role.max for count in Counter(raw_memberships).values() + ) + + group_values = np.arange(group.count, dtype=float) + group_indices = {group_id: index for index, group_id in enumerate(group.ids)} + counts = np.array( + [(raw_memberships == group_id).sum() for group_id in group.ids] + ) + totals = np.array( + [values[raw_memberships == group_id].sum() for group_id in group.ids] + ) + projected_values = np.array( + [group_values[group_indices[group_id]] for group_id in raw_memberships] + ) + + for query_role in [*entity.roles, *entity.flattened_roles]: + holds = role is not None and ( + query_role is role or role in (query_role.subroles or []) + ) + has_role = simulation.persons.has_role(query_role) + nb_persons = group.nb_persons(role=query_role) + total = group.sum(values, role=query_role) + projected = group.project(group_values, role=query_role) + + assert has_role.shape == projected.shape == (persons,) + assert nb_persons.shape == total.shape == (group.count,) + assert np.array_equal(has_role, np.full(persons, holds)) + assert np.array_equal( + nb_persons, counts if holds else np.zeros(group.count) + ) + if query_role.max is not None: + assert (nb_persons <= query_role.max).all() + assert np.array_equal(total, totals if holds else np.zeros(group.count)) + assert np.array_equal( + projected, projected_values if holds else np.zeros(persons) + ) + if persons: + maximum = group.max(values, role=query_role) + expected_maximum = np.array( + [ + max(values[raw_memberships == group_id], default=-np.inf) + if holds + else -np.inf + for group_id in group.ids + ] + ) + assert maximum.shape == (group.count,) + assert np.array_equal(maximum, expected_maximum) + + for unique_role in entity.flattened_roles: + if unique_role.max != 1: + continue + # The accepted-capacity invariant guarantees zero or one selected + # person per raw group ID. Unmatched roles select nobody. + expected_lookup = np.array( + [ + next(iter(values[raw_memberships == group_id]), -123) + if role is unique_role + else -123 + for group_id in group.ids + ] + ) + lookup = group.value_from_person(values, unique_role, default=-123) + assert lookup.shape == (group.count,) + assert np.array_equal(lookup, expected_lookup) + + expected_projector = np.array( + [ + next(iter(values[raw_memberships == group_id]), 0) + if role is unique_role + else 0 + for group_id in group.ids + ] + ) + projector = UniqueRoleToEntityProjector(group, unique_role) + assert np.array_equal(projector.transform(values), expected_projector) + + for compound_role in entity.roles: + if not compound_role.subroles or len(compound_role.subroles) != 2: + continue + # Uniform defaults cannot supply both partner subroles. Pass the + # public person-to-group projector so group values project back to + # person shape, as in existing partner contract tests. + partners = simulation.persons.value_from_partner( + values, getattr(simulation.persons, entity.key), compound_role + ) + assert partners.shape == (persons,) + assert not partners.any() + + +@settings(max_examples=50, deadline=None) +@given( + st.integers(min_value=1, max_value=50), st.sampled_from(["short", "long", "matrix"]) +) +@example(0, "matrix") +def test_join_requires_a_one_dimensional_role_per_person(persons, shape): + builder = SimulationBuilder() + builder.create_entities(SYSTEM) + builder.declare_person_entity("person", np.arange(persons)) + household = builder.declare_entity("household", [10]) + if shape == "matrix": + roles = np.zeros((persons, 1), dtype=int) + else: + roles = np.zeros(persons + (1 if shape == "long" else -1), dtype=int) + + with pytest.raises(ValueError, match="give one role per person"): + builder.join_with_persons(household, np.repeat(10, persons), roles) + + +@settings(max_examples=100, deadline=None) +@given( + st.lists(st.integers(min_value=0, max_value=2), min_size=1, max_size=12), + st.booleans(), + st.booleans(), + st.sampled_from(["2024", "eternity"]), +) +@example([0, 2, 2], False, False, "2024") +@example([0, 0, 0], True, True, "eternity") +def test_suffixed_entity_roles_preserve_explicit_assignment( + role_indices, generic_role, literal_keys, suffix +): + # Independent encoding table: do not derive expected roles from the loader. + role_table = { + 0: ("first_parent", HOUSEHOLD.FIRST_PARENT), + 1: ("second_parent", HOUSEHOLD.SECOND_PARENT), + 2: ("child", HOUSEHOLD.CHILD), + } + persons = len(role_indices) + explicit_roles = ( + [role_table[index][0] for index in role_indices] + if literal_keys + else role_indices + ) + columns = { + "person_id": np.arange(1, persons + 1), + # Group zero matches the flat-file loader's synthetic group IDs. + "person_household_id": np.zeros(persons, dtype=int), + f"person_household_role__{suffix}": explicit_roles, + } + if generic_role: + # Every generic entry conflicts with its entity-specific counterpart. + columns["role"] = [(index + 1) % 3 for index in role_indices] + simulation = simulate( + TaxBenefitSystem([PERSON, HOUSEHOLD]), + Dataset.from_dataframe(pd.DataFrame(columns), "2024"), + default_role="first_parent", + ) + household = simulation.household + assert list(household.members_role) == [ + role_table[index][1] for index in role_indices + ] + + # Explicit columns retain their existing capacity handling: repeated unique + # roles remain accepted. Counts and sums still follow the declared roles. + values = np.arange(1, persons + 1) + query_roles = [ + (HOUSEHOLD.FIRST_PARENT, {0}), + (HOUSEHOLD.SECOND_PARENT, {1}), + (HOUSEHOLD.CHILD, {2}), + (HOUSEHOLD.PARENT, {0, 1}), + ] + for role, matching_indices in query_roles: + expected_count = sum(index in matching_indices for index in role_indices) + expected_sum = sum( + value + for value, index in zip(values, role_indices) + if index in matching_indices + ) + assert household.nb_persons(role=role).tolist() == [expected_count] + assert household.sum(values, role=role).tolist() == [expected_sum]