Keep loaded roles and IDs through subsample without backing variables

[MaxGhenis]

Draft until the independent review comes back. The fix, the example tests and the property tests are in.

Problem

Simulation.subsample() rebuilds the simulation through build_from_dataset from a person-level frame. That frame holds the recorded inputs plus the structural leaves that retain_structure adds. retain_structure skipped any leaf the tax-benefit system has no variable for. The loader reads those columns by name, not through a variable, so whatever they held was lost:

• Roles. With no person_<entity>_role variable, the rebuild fell back to default_role, so every kept member got the default role.
  • On master this is silent.
  • With Validate implicit dataset role capacities and role array shapes #595 merged, a bounded default raises its new capacity ValueError part-way through the rebuild. That leaves the simulation unusable (members_entity_id is None).
  • With default_role=None, master raised Missing person_household_role column.
• IDs. With no person_id variable, subsample raised KeyError part-way through the rebuild. With no person_<entity>_id variable, it raised IndexError.

Found in review P3-1 on #595.

No shipped country package reaches the silent case today:

• US has no role variables and every US role is member, which is also its default.
• UK's role variables are strings with recorded inputs.
• CA, IL and NG default to "parent", which matches no role, and have no role variables.

The subsample path still threw away roles the caller supplied.

Before and after

The reviewer's case: the country template with person_household_role removed and default_role = "first_parent". The flat file has person_household_id__2024 = [0, 0, 0, 1] and person_household_role__2024 = [first_parent, child, child, first_parent]. Then subsample(n=2, seed=0):

	first_parent per household	child per household
Build (all commits)	[1, 1]	[2, 0]
subsample, master 1cccc278e	[3, 1] (everyone first_parent)	[0, 0]
subsample, master + #595 ef0f9d5e2c	raises ValueError: Cannot assign default role 'first_parent'…; simulation left unusable	n/a
subsample, this PR	[1, 1]	[2, 0]

Changes

All in policyengine_core/simulations/simulation.py.

• retain_structure (in subsample) now writes a structural leaf whether or not a variable backs it.
  • A leaf with no variable is written as <name>__ETERNITY.
  • A leaf that is already a recorded input is left alone, as before.
• Roles are carried by key. The key comes from each member's members_role. A member the loader matched to no role holds 0; it is written as "" and reloads as no role.
• Integer role variables.
  • The column holds flattened-role indices, now looked up by key rather than by object identity.
  • An integer column cannot say "no role". If any member has no role, subsample raises ValueError before changing anything. Master wrote index 0 here, which silently gave those members the first flattened role.
• build_from_dataset records each column name it reads for structure (through get_eternity_array). It no longer reports those columns as "ignored" in the unknown-column warning, because they are used. Without this, a US subsample would warn about the five member role columns it now carries. The original load already warned, wrongly, that an explicit role column with no variable was ignored. Genuinely unknown columns still warn.

Invariants

These hold for every simulation loaded from a dataset and every subsample(n, seed, quantize_weights):

1. Roles kept. For every kept person, the role after subsampling equals the role it was loaded with. "No role" stays no role.
2. Whole households. Every kept household is a whole source household: its member set equals one source household's member set.
3. Role counts kept. For every kept household, the count of each role key equals the source household's count, and nb_persons(role) agrees with those counts for every flattened role.
4. Independent of variables and defaults. Invariants 1–3 hold whether or not the system defines person_household_role, whatever default_role is (bounded, unbounded, unmatched or None), and whether roles came from the entity column or the generic role column.
5. Repeatable. Subsampling again preserves invariants 1–3.
6. Fail before changing. With an integer role variable, subsample either keeps roles (1–3) or raises ValueError. It raises exactly when some member has no flattened role, and then the simulation is unchanged: same dataset object, persons, roles and weights.
7. Warnings. No unknown-column warning names a column the loader reads for structure.

Tests

The oracle is a hand-written household table in tests/fixtures/subsample_roles.py; expected roles are never taken from the loader.

• tests/core/test_subsample_explicit_roles.py: 31 example tests.

  • With the new tests copied onto master 1cccc278e, 21 fail for the reasons above.
  • 10 pass on master. These lock in behaviour that already works: the 8 cases with a string role variable (master's existing path), an integer variable whose members all hold flattened roles, and a subsample with no role column at all.

• tests/core/test_subsample_explicit_roles_property.py: two Hypothesis properties. Together they cover invariants 1–7.

  • Households mix bounded, unbounded and unmatched keys. Person order, weights, default_role, the role source, n, the seed and quantize_weights all vary.
  • Results: 200 and 100 passing examples on this branch; both properties fail on master.
  • On master, generation alone (with the explicit example removed) finds several distinct failures. The integer property shrinks to a single lodger member, where master writes role index 0 instead of raising.

• Master merged with Validate implicit dataset role capacities and role array shapes #595 (ef0f9d5e2c; local merge, clean): the reviewer's case keeps first_parent [1, 1] and child [2, 0], with no capacity error. These files pass one at a time:

  File	Result
  test_subsample_explicit_roles.py	31 passed
  test_subsample_explicit_roles_property.py	2 passed
  test_dataset_default_role.py	35 passed
  test_dataset_default_role_property.py	3 passed
  test_subsample_structural_ids.py	27 passed
  test_dataset_unknown_column_warning.py	2 passed
  test_simulation_builder.py	35 passed

• Existing files, run one at a time on this branch: all pass.

  • All 14 tests/core/test_subsample_*.py files, including the three existing property modules.
  • test_dataset_unknown_column_warning.py, test_dataset_getattr.py, test_simulation_builder.py, test_dump_restore.py, test_bincount_minlength.py, test_entities.py, test_projectors.py, test_microsimulation_weights.py and test_microsimulation_person_accessor.py.
  • Warning counts in the subsample files match master.

• Lint: ruff format --check . and ruff check are clean.

• Cost: reading role keys takes 48 ms for 5 group entities × 400k persons, about the size of a full US subsample.

Not run here: the PE-US core gate and full-sample controls. The US + core hub runs those.

axiom: n/a: core engine infrastructure, no policy change.

🤖 Generated with Claude Code