Keep loaded roles and IDs through subsample without backing variables

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.

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

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_<entity>_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()