diff --git a/changelog.d/subsample-preserve-roles.fixed.md b/changelog.d/subsample-preserve-roles.fixed.md new file mode 100644 index 000000000..bc846b91b --- /dev/null +++ b/changelog.d/subsample-preserve-roles.fixed.md @@ -0,0 +1 @@ +`Simulation.subsample` now keeps each member's loaded role (and person and membership IDs) when the tax-benefit system has no variable for them, instead of rebuilding everyone with the default role, and the dataset loader no longer reports the structural columns it reads as ignored. diff --git a/policyengine_core/simulations/simulation.py b/policyengine_core/simulations/simulation.py index ddda2ab94..676136b1e 100644 --- a/policyengine_core/simulations/simulation.py +++ b/policyengine_core/simulations/simulation.py @@ -1101,8 +1101,12 @@ def build_from_dataset(self) -> None: person_entity = self.tax_benefit_system.person_entity entity_id_field = f"{person_entity.key}_id" + # Columns read to build the populations are used even when no + # variable has their name, so they are not reported as ignored. + structural_columns = set() def get_eternity_array(name): + structural_columns.add(name) if self.dataset.data_format == Dataset.FLAT_FILE: # Look for any column with variablename__timeperiod for col in data: @@ -1187,7 +1191,7 @@ def get_eternity_array(name): self.set_input( variable, self.dataset.time_period, data[variable] ) - else: + elif variable not in structural_columns: unknown_columns.append(variable) else: for variable in data: @@ -1198,7 +1202,8 @@ def get_eternity_array(name): time_period = self.dataset.time_period or self.default_input_period if variable_name not in self.tax_benefit_system.variables: - unknown_columns.append(variable) + if variable_name not in structural_columns: + unknown_columns.append(variable) continue variable_meta = self.tax_benefit_system.get_variable(variable_name) @@ -3375,6 +3380,11 @@ def subsample( If structural inputs are absent, loaded population IDs, memberships and roles supply the leaves needed to rebuild the same partition. This also supports structural variables whose IDs have formulas. + These leaves are carried whether or not the system defines a variable + for them, and roles are carried by key, so every kept member keeps + the role it was loaded with instead of the default role. An integer + role variable cannot represent a member with no role; subsampling + raises ``ValueError`` before changing anything in that case. Args: n (int, optional): The number of households to sample. Defaults to 10_000. @@ -3404,12 +3414,15 @@ def subsample( df = self._to_person_dataframe(self._get_set_input_periods) def retain_structure(name, values): - variable = self.tax_benefit_system.get_variable(name) - if variable is None or any(column.split("__")[0] == name for column in df): + if any(column.split("__")[0] == name for column in df): return + # The loader reads structural columns by name, so a leaf that no + # variable backs (a role column the system does not define, for + # instance) is still carried. It has no period of its own. + variable = self.tax_benefit_system.get_variable(name) structural_period = ( ETERNITY - if variable.definition_period == ETERNITY + if variable is None or variable.definition_period == ETERNITY else periods.period(self.dataset.time_period).start.period( variable.definition_period ) @@ -3427,18 +3440,35 @@ def retain_structure(name, values): population.ids[population.members_entity_id], ) role_name = f"{person_key}_{entity.key}_role" + if any(column.split("__")[0] == role_name for column in df): + continue + # Carry each member's own role by key; the rebuild otherwise + # gives everyone the default. Members the loader matched to no + # role hold 0, and an empty key reloads them the same way. role_variable = self.tax_benefit_system.get_variable(role_name) - roles = np.broadcast_to(population.members_role, self.persons.count) + roles = np.broadcast_to( + np.asarray(population.members_role, dtype=object), + self.persons.count, + ) + role_keys = [getattr(role, "key", "") for role in roles] if role_variable is not None and role_variable.value_type == int: - role_values = [ - entity.flattened_roles.index(role) - if role in entity.flattened_roles - else 0 - for role in roles - ] + # An integer role column indexes flattened roles, so it has + # no value for "no role". Stop before anything is replaced. + index_of = { + role.key: index for index, role in enumerate(entity.flattened_roles) + } + unindexed = sorted({key for key in role_keys if key not in index_of}) + if unindexed: + raise ValueError( + f"Cannot subsample: {role_name} is an integer role " + f"variable, so it can only hold the {entity.key} roles " + f"{[role.key for role in entity.flattened_roles]}, but " + "some members hold " + f"{[key or '(no role)' for key in unindexed]}." + ) + retain_structure(role_name, [index_of[key] for key in role_keys]) else: - role_values = [getattr(role, "key", "") for role in roles] - retain_structure(role_name, role_values) + retain_structure(role_name, role_keys) # Flat-file loading can infer households from membership IDs even # when the household ID variable has a formula. Use the same leaves diff --git a/tests/core/test_subsample_explicit_roles.py b/tests/core/test_subsample_explicit_roles.py new file mode 100644 index 000000000..627253f91 --- /dev/null +++ b/tests/core/test_subsample_explicit_roles.py @@ -0,0 +1,240 @@ +"""``subsample`` keeps the roles members were loaded with. + +The rebuild reads each member's role from a ``person__role`` column. +That column used to be carried only when the system defined a variable of +that name, so a system without one rebuilt every member with the default +role: silently on master, and as a capacity error once implicit defaults +are capacity-checked (review P3-1 on PolicyEngine/policyengine-core#595). +Person and membership IDs had the same gap. +""" + +import logging + +import pandas as pd +import pytest + +from policyengine_core.data import Dataset +from policyengine_core.periods import ETERNITY +from tests.fixtures.subsample_roles import ( + DATASET_YEAR, + NO_ROLE, + ROLE, + UNKNOWN_COLUMN_WARNING, + assert_roles_kept, + load, + load_dataframe, + role_counts, + role_keys, + role_system, + roles_by_person, + source_roles, +) + +HOUSEHOLDS = [ + ["first_parent", "second_parent", "child", "child"], + ["first_parent", "child"], + ["first_parent", "lodger"], + ["second_parent", "first_parent", "child", "parent"], + ["child"], +] +# Shuffled so that a household's members are not adjacent rows. +ORDER = [7, 0, 12, 3, 9, 1, 5, 10, 2, 8, 11, 4, 6] + + +def test_review_reproduction_keeps_explicit_roles(): + """The reviewed case: no role variable and a bounded default role.""" + data = pd.DataFrame( + { + f"person_id__{DATASET_YEAR}": [0, 1, 2, 3], + f"person_household_id__{DATASET_YEAR}": [0, 0, 0, 1], + f"{ROLE}__{DATASET_YEAR}": [ + "first_parent", + "child", + "child", + "first_parent", + ], + f"household_weight__{DATASET_YEAR}": [1.0, 1.0, 1.0, 1.0], + } + ) + simulation = load_dataframe(data, default_role="first_parent") + assert role_counts(simulation) == { + "first_parent": [1, 1], + "second_parent": [0, 0], + "child": [2, 0], + } + + simulation.subsample(n=2, seed=0) + + # Master rebuilt everyone as first_parent: first_parent [3, 1], child [0, 0]. + assert simulation.household.count == 2 + assert role_counts(simulation) == { + "first_parent": [1, 1], + "second_parent": [0, 0], + "child": [2, 0], + } + assert roles_by_person(simulation) == { + 0: "first_parent", + 1: "child", + 2: "child", + 3: "first_parent", + } + + +@pytest.mark.parametrize("role_variable", ["absent", "string"]) +@pytest.mark.parametrize("role_column", [ROLE, "role"]) +@pytest.mark.parametrize("default_role", ["first_parent", "child", "parent", None]) +def test_explicit_roles_survive_subsample(role_variable, role_column, default_role): + simulation = load( + HOUSEHOLDS, + role_variable=role_variable, + role_column=role_column, + default_role=default_role, + order=ORDER, + ) + assert_roles_kept(simulation, HOUSEHOLDS) + assert simulation.household.count == len(HOUSEHOLDS) + + simulation.subsample(n=8, seed="roles", time_period=DATASET_YEAR) + + assert simulation.household.count >= 2 + assert_roles_kept(simulation, HOUSEHOLDS) + + +def test_roles_survive_repeated_subsamples(): + simulation = load(HOUSEHOLDS, order=ORDER, default_role="child") + for n, seed in [(8, "first"), (3, "second"), (2, "third")]: + simulation.subsample(n=n, seed=seed, time_period=DATASET_YEAR) + assert_roles_kept(simulation, HOUSEHOLDS) + + +def test_rebuilt_dataset_carries_role_keys(): + simulation = load(HOUSEHOLDS, order=ORDER) + simulation.subsample(n=8, seed="roles", time_period=DATASET_YEAR) + + data = simulation.dataset.load() + [role_column] = [column for column in data if column.split("__")[0] == ROLE] + [id_column] = [column for column in data if column.split("__")[0] == "person_id"] + assert role_column == f"{ROLE}__{ETERNITY}" + expected = source_roles(HOUSEHOLDS) + people = data[id_column].tolist() + assert dict(zip(people, data[role_column].tolist())) == { + person: expected[person] for person in people + } + + +@pytest.mark.parametrize("missing", ["person_id", "person_household_id"]) +def test_structural_ids_without_variables_survive_subsample(missing): + """Master raised: ``KeyError`` mid-rebuild, or ``IndexError``.""" + system = role_system("absent") + del system.variables[missing] + simulation = load(HOUSEHOLDS, order=ORDER, system=system) + source_ids = sorted(simulation.persons.ids.tolist()) + + simulation.subsample(n=8, seed="roles", time_period=DATASET_YEAR) + + assert set(simulation.persons.ids.tolist()) <= set(source_ids) + assert_roles_kept(simulation, HOUSEHOLDS) + + +@pytest.mark.parametrize("unmatched", ["lodger", "parent"]) +def test_integer_role_variable_rejects_roleless_members_before_changing_anything( + unmatched, +): + """An integer role column cannot say "no role"; master wrote role 0.""" + households = [["first_parent", "child"], ["first_parent", unmatched]] + simulation = load(households, role_variable="integer", role_column="role") + dataset = simulation.dataset + persons = simulation.persons.ids.tolist() + roles = role_keys(simulation).tolist() + weights = simulation.calculate("household_weight", DATASET_YEAR).tolist() + + with pytest.raises(ValueError, match="integer role variable.*no role"): + simulation.subsample(n=2, seed="roles", time_period=DATASET_YEAR) + + assert simulation.dataset is dataset + assert simulation.persons.ids.tolist() == persons + assert role_keys(simulation).tolist() == roles + assert simulation.calculate("household_weight", DATASET_YEAR).tolist() == weights + + +def test_integer_role_variable_carries_flattened_roles(): + households = [["first_parent", "second_parent", "child"], ["child", "first_parent"]] + simulation = load(households, role_variable="integer", role_column="role") + + simulation.subsample(n=4, seed="roles", time_period=DATASET_YEAR) + + assert_roles_kept(simulation, households) + + +def unknown_column_warnings(caplog): + return [ + record.getMessage() + for record in caplog.records + if UNKNOWN_COLUMN_WARNING in record.getMessage() + ] + + +@pytest.mark.parametrize("role_column", [ROLE, "role"]) +def test_structural_columns_are_not_reported_as_ignored(caplog, role_column): + """The loader reads these columns, though no variable has their name.""" + system = role_system("absent") + del system.variables["person_household_id"] + with caplog.at_level(logging.WARNING): + simulation = load(HOUSEHOLDS, role_column=role_column, system=system) + + assert unknown_column_warnings(caplog) == [] + assert_roles_kept(simulation, HOUSEHOLDS) + + +@pytest.mark.parametrize("role_column", [ROLE, "role", None]) +def test_carried_role_columns_are_not_reported_as_ignored(caplog, role_column): + """Like a system with no role variables (policyengine-us, for instance), + whose subsample now carries a role column per group entity.""" + households = HOUSEHOLDS if role_column else [["child"] * 3, ["child"], ["child"]] + simulation = load(households, role_column=role_column, default_role="child") + caplog.clear() + with caplog.at_level(logging.WARNING): + simulation.subsample(n=8, seed="roles", time_period=DATASET_YEAR) + + assert unknown_column_warnings(caplog) == [] + assert_roles_kept(simulation, households) + + +def test_unknown_columns_are_still_reported_beside_structural_ones(caplog): + data = load(HOUSEHOLDS).dataset.load() + data[f"a_removed_input__{DATASET_YEAR}"] = 1.0 + with caplog.at_level(logging.WARNING): + load_dataframe(data) + + [warning] = unknown_column_warnings(caplog) + assert "1 column(s)" in warning + assert f"a_removed_input__{DATASET_YEAR}" in warning + assert ROLE not in warning + + +def test_time_period_array_role_column_is_used_and_not_reported(caplog): + frame = load(HOUSEHOLDS).dataset.load() + data = {column.split("__")[0]: {ETERNITY: frame[column].values} for column in frame} + data["household_id"] = {ETERNITY: [10, 20, 30, 40, 50]} + data["household_weight"] = {DATASET_YEAR: [1.0, 2.0, 3.0, 4.0, 5.0]} + dataset = type( + "RoleArrays", + (Dataset,), + { + "name": "role_arrays", + "label": "Role arrays", + "data_format": Dataset.TIME_PERIOD_ARRAYS, + "file_path": "dataframe", + "time_period": DATASET_YEAR, + "load": lambda self: data, + }, + )() + simulation_class = type(load(HOUSEHOLDS)) + with caplog.at_level(logging.WARNING): + simulation = simulation_class( + tax_benefit_system=role_system("absent"), dataset=dataset + ) + + assert unknown_column_warnings(caplog) == [] + assert_roles_kept(simulation, HOUSEHOLDS) + assert NO_ROLE in role_keys(simulation).tolist() diff --git a/tests/core/test_subsample_explicit_roles_property.py b/tests/core/test_subsample_explicit_roles_property.py new file mode 100644 index 000000000..60728a64d --- /dev/null +++ b/tests/core/test_subsample_explicit_roles_property.py @@ -0,0 +1,183 @@ +"""Property: ``subsample`` keeps every kept member's loaded role. + +For any households of explicit role keys (bounded, unbounded and unmatched), +any person order, weights, default role, role source and sample, every kept +household is a whole source household whose role counts, by key and by +``nb_persons``, equal the source household's, whether or not the system +defines a role variable. An integer role variable either keeps the roles or +rejects members with no role before changing anything. +""" + +import logging +from contextlib import contextmanager + +import pytest + +from tests.fixtures.subsample_roles import ( + DATASET_YEAR, + FLATTENED_KEYS, + ROLE, + UNKNOWN_COLUMN_WARNING, + assert_roles_kept, + load, + role_keys, +) + +hypothesis = pytest.importorskip("hypothesis") +st = hypothesis.strategies + +# ``parent`` is the template's compound key and ``lodger`` no key at all: +# both load as no role. +EXTRA_KEYS = ["child", "lodger", "parent"] + + +def household(first_parent, second_parent): + parents = ["first_parent"] * first_parent + ["second_parent"] * second_parent + return st.lists( + st.sampled_from(EXTRA_KEYS), min_size=0 if parents else 1, max_size=3 + ).map(lambda others: parents + others) + + +HOUSEHOLD = st.tuples(st.booleans(), st.booleans()).flatmap( + lambda parents: household(*parents) +) +HOUSEHOLDS = st.lists(HOUSEHOLD, min_size=1, max_size=6) +DEFAULT_ROLES = st.sampled_from( + ["first_parent", "second_parent", "child", "parent", None] +) +SEEDS = st.sampled_from([0, 1, "a", "b"]) + + +@contextmanager +def unknown_column_warnings(): + class Collect(logging.Handler): + def __init__(self): + super().__init__(logging.WARNING) + self.messages = [] + + def emit(self, record): + if UNKNOWN_COLUMN_WARNING in record.getMessage(): + self.messages.append(record.getMessage()) + + handler = Collect() + root = logging.getLogger() + root.addHandler(handler) + try: + yield handler.messages + finally: + root.removeHandler(handler) + + +def draw_layout(data, households): + people = sum(len(household) for household in households) + order = data.draw(st.permutations(range(people)), label="order") + weights = data.draw( + st.lists(st.integers(1, 9), min_size=len(households), max_size=len(households)), + label="weights", + ) + return order, [float(weight) for weight in weights] + + +@hypothesis.settings(max_examples=200, deadline=None) +@hypothesis.example( + households=[["first_parent", "child", "child"], ["first_parent"]], + role_variable="absent", + role_column=ROLE, + default_role="first_parent", + n=2, + seed=0, + quantize_weights=True, + again=1, + data=None, +).via("review P3-1 on #595") +@hypothesis.given( + households=HOUSEHOLDS, + role_variable=st.sampled_from(["absent", "string"]), + role_column=st.sampled_from([ROLE, "role"]), + default_role=DEFAULT_ROLES, + n=st.integers(1, 8), + seed=SEEDS, + quantize_weights=st.booleans(), + again=st.integers(1, 4), + data=st.data(), +) +def test_subsample_keeps_explicit_roles( + households, + role_variable, + role_column, + default_role, + n, + seed, + quantize_weights, + again, + data, +): + order, weights = (None, None) if data is None else draw_layout(data, households) + with unknown_column_warnings() as warnings: + simulation = load( + households, + role_variable=role_variable, + role_column=role_column, + default_role=default_role, + order=order, + weights=weights, + ) + # The loader itself agrees with the table, so later checks test + # what subsample keeps rather than what was loaded. + assert_roles_kept(simulation, households) + assert simulation.household.count == len(households) + + simulation.subsample( + n=n, seed=seed, time_period=DATASET_YEAR, quantize_weights=quantize_weights + ) + assert 1 <= simulation.household.count <= min(n, len(households)) + assert_roles_kept(simulation, households) + + # The rebuilt simulation carries them on through another subsample. + simulation.subsample(n=again, seed=seed, time_period=DATASET_YEAR) + assert_roles_kept(simulation, households) + + assert warnings == [] + + +@hypothesis.settings(max_examples=100, deadline=None) +@hypothesis.given( + households=HOUSEHOLDS, + default_role=DEFAULT_ROLES, + n=st.integers(1, 8), + seed=SEEDS, + data=st.data(), +) +def test_integer_role_variable_keeps_roles_or_changes_nothing( + households, default_role, n, seed, data +): + order, weights = draw_layout(data, households) + simulation = load( + households, + role_variable="integer", + role_column="role", + default_role=default_role, + order=order, + weights=weights, + ) + roleless = any( + key not in FLATTENED_KEYS for household in households for key in household + ) + dataset = simulation.dataset + persons = simulation.persons.ids.tolist() + roles = role_keys(simulation).tolist() + household_weights = simulation.calculate("household_weight", DATASET_YEAR).tolist() + + if roleless: + with pytest.raises(ValueError, match="integer role variable"): + simulation.subsample(n=n, seed=seed, time_period=DATASET_YEAR) + assert simulation.dataset is dataset + assert simulation.persons.ids.tolist() == persons + assert role_keys(simulation).tolist() == roles + assert ( + simulation.calculate("household_weight", DATASET_YEAR).tolist() + == household_weights + ) + else: + simulation.subsample(n=n, seed=seed, time_period=DATASET_YEAR) + assert_roles_kept(simulation, households) diff --git a/tests/fixtures/subsample_roles.py b/tests/fixtures/subsample_roles.py new file mode 100644 index 000000000..142014a4e --- /dev/null +++ b/tests/fixtures/subsample_roles.py @@ -0,0 +1,185 @@ +"""Flat-file households with explicit roles, for testing what ``subsample`` keeps. + +Shared by ``tests/core/test_subsample_explicit_roles.py`` and its Hypothesis +property module, which pytest must be able to skip on its own when +Hypothesis is not installed. + +Households are written as tables of role keys, and the expected role of each +person is read from that table, never from the loader: a key that names a +flattened household role (``first_parent``, ``second_parent``, ``child``) is +that role, and any other key (``parent``, ``lodger``) is no role. +""" + +from collections import Counter + +import numpy as np +import pandas as pd + +from policyengine_core.country_template import CountryTaxBenefitSystem, Simulation +from policyengine_core.country_template.entities import Person +from policyengine_core.data import Dataset +from policyengine_core.periods import ETERNITY +from policyengine_core.variables import Variable + +DATASET_YEAR = "2024" +ROLE = "person_household_role" +FLATTENED_KEYS = ("first_parent", "second_parent", "child") +# How these helpers report a member the loader gave no role. +NO_ROLE = "" +UNKNOWN_COLUMN_WARNING = "do not match any variable" + + +class integer_role(Variable): + value_type = int + entity = Person + definition_period = ETERNITY + label = "Integer person household role" + + def formula(person, period): + return person.filled_array(2) + + +integer_role.__name__ = ROLE + + +def role_system(role_variable: str = "absent") -> CountryTaxBenefitSystem: + """The template system with its household role variable ``absent``, + as shipped (``string``), or replaced by an ``integer`` formula.""" + system = CountryTaxBenefitSystem() + if role_variable == "absent": + del system.variables[ROLE] + elif role_variable == "integer": + system.replace_variable(integer_role) + else: + assert role_variable == "string" + return system + + +def expected_role(key: str) -> str: + return key if key in FLATTENED_KEYS else NO_ROLE + + +def dataframe(households, role_column=ROLE, order=None, weights=None): + """One row per person: person ``i`` has ID ``100 + i``, household ``h`` + has membership label ``10 * (h + 1)`` and weight ``weights[h]``.""" + keys = [key for household in households for key in household] + labels = [10 * (h + 1) for h, household in enumerate(households) for _ in household] + weights = weights or [float(h + 1) for h in range(len(households))] + columns = { + f"person_id__{DATASET_YEAR}": [100 + i for i in range(len(keys))], + f"person_household_id__{DATASET_YEAR}": labels, + f"household_weight__{DATASET_YEAR}": [ + float(weights[label // 10 - 1]) for label in labels + ], + } + if role_column == ROLE: + columns[f"{ROLE}__{DATASET_YEAR}"] = keys + elif role_column is not None: + columns[role_column] = keys + data = pd.DataFrame(columns) + if order is not None: + data = data.iloc[list(order)].reset_index(drop=True) + return data + + +def source_roles(households) -> dict: + """``{person ID: expected role key}`` from the household table.""" + keys = [key for household in households for key in household] + return {100 + i: expected_role(key) for i, key in enumerate(keys)} + + +def source_households(households) -> dict: + """``{frozenset of person IDs: Counter of expected role keys}``.""" + result, person = {}, 100 + for household in households: + ids = frozenset(range(person, person + len(household))) + result[ids] = Counter(expected_role(key) for key in household) + person += len(household) + return result + + +def load_dataframe( + data, role_variable="absent", default_role="first_parent", system=None +): + simulation_class = type( + "RoleSimulation", (Simulation,), {"default_role": default_role} + ) + return simulation_class( + tax_benefit_system=system or role_system(role_variable), + dataset=Dataset.from_dataframe(data, DATASET_YEAR), + ) + + +def load( + households, + role_variable="absent", + role_column=ROLE, + default_role="first_parent", + order=None, + weights=None, + system=None, +): + return load_dataframe( + dataframe(households, role_column, order, weights), + role_variable, + default_role, + system, + ) + + +def role_keys(simulation) -> np.ndarray: + household = simulation.populations["household"] + roles = np.broadcast_to( + np.asarray(household.members_role, dtype=object), simulation.persons.count + ) + return np.array([getattr(role, "key", NO_ROLE) for role in roles], dtype=object) + + +def roles_by_person(simulation) -> dict: + return dict(zip(simulation.persons.ids.tolist(), role_keys(simulation).tolist())) + + +def households_by_members(simulation) -> dict: + """``{frozenset of person IDs: Counter of role keys}`` per household.""" + household = simulation.populations["household"] + keys = role_keys(simulation) + result = {} + for index in range(household.count): + members = household.members_entity_id == index + ids = frozenset(simulation.persons.ids[members].tolist()) + result[ids] = Counter(keys[members].tolist()) + return result + + +def role_counts(simulation) -> dict: + """``{role key: nb_persons(role) per household}``, from the projector.""" + household = simulation.populations["household"] + return { + role.key: household.nb_persons(role).tolist() + for role in household.entity.flattened_roles + } + + +def assert_roles_kept(simulation, households): + """Every kept household is a whole source household with its roles.""" + source = source_households(households) + kept = households_by_members(simulation) + assert kept, "subsample kept no household" + for members, roles in kept.items(): + assert members in source, f"{sorted(members)} is not a source household" + assert roles == source[members] + expected = source_roles(households) + assert roles_by_person(simulation) == { + person: expected[person] for person in simulation.persons.ids.tolist() + } + # The projector counts the same roles the members hold. + household_order = [ + frozenset( + simulation.persons.ids[ + simulation.household.members_entity_id == index + ].tolist() + ) + for index in range(simulation.household.count) + ] + for key, counts in role_counts(simulation).items(): + assert counts == [kept[members][key] for members in household_order]