[MaxGhenis]

Follow-up to #590. Its reviews found two neighbouring defects outside its scope. This PR fixes both. It leaves #590's behaviour for integer counts unchanged.

What this fixes

1. Restoring a dump with no persons, or with groups that have no members

restore_simulation set each group's count to one past its highest member index, and the number of persons to the length of a group's memberships.

• No persons. That was max() of an empty array: ValueError: max() iterable argument is empty. This predates Reduce and rank over a population with no persons #590.
• Empty groups last. With five households and members only in the second and fourth, the count came out as 4, not 5. Restoring any household value then raised ValueError: Unable to set value ... as its length is 5 while there are 4 households. Empty groups first or in between were restored correctly.
• No group entities. The number of persons was never set: UnboundLocalError.

dump_simulation now writes each population's count, as an int64 scalar, to __entities__/<entity>/count.npy. restore_simulation restores it.

A population with no count.npy gets the smallest count its ids and memberships allow:

• at least one per id;
• for a group, one past its highest member index;
• for persons, as many as the largest group membership list.

That covers two cases:

• Dumps written before this change. For a group, the inferred count differs from the old one only when the group has more ids than one past its highest member index, that is, when its last groups have no members.
• A count that is not an integer. Only code that sets the count by hand can make one; the dump writes no file for it, so dumping it still works.

Older Core versions load only the files they know by name, so they still restore new dumps as before.

2. Non-integer counts in build_default_simulation

Since #590, build_default_simulation uses np.arange(count, dtype=np.int64). That accepts 2.5, 2.0, np.float64(2.0), np.float32, 0-d float arrays, np.True_, Decimal and Fraction. The simulation then fails later:

• with 2.5, sum raises size = 2 != 2.5 = count;
• with 2.0, set_input works, but sum and calculate raise TypeError.

Before #590, range(count) raised TypeError at build. The method now calls operator.index(count) first. That refuses exactly what range refuses, because range also converts its argument with __index__.

For Python and NumPy integers and 0-d integer arrays, population.count and the int64 ids are unchanged from #590. An object that is an integer only through __index__ builds again, as it did before #590; np.arange refused it. set_input then refuses every array for such a count, because holders compare an array's length with the count.

No caller in the repo passes a float; _get_person_count returns len(...) or 1.

Invariants

These are the properties the code is meant to hold. The property tests check them by sampling, over the domain they name.

• Round trip. Take a simulation in which:

  • every population's count is an integer, of any integer type;
  • its ids are numeric or fixed-width strings, or were never set. Object-dtype ids can't be restored, at master either.

  Dump it, then restore it. Then:

  • every population has the same count and ids;
  • every group has the same memberships, positions and roles;
  • every holder has the same values, with the same array type, enum class, dtype, shape, bytes and derived mark;
  • a calculation gives the same result or error.

  This includes zero persons, zero groups, and groups with no members first, last or in between. Every count file is int64.

  A count that is not an integer can only be set by hand, and the dump writes none for it. Restoring infers it under the next rule. So it comes back as the integer the ids or memberships give, which can differ in behaviour: with a hand-set 2.0, sum raises before the dump and works after. With no ids and an empty last group, the inferred count is too small and restoring that group's values raises, as at master. test_dump_restore_property.py samples it on the country template:

  • up to 6 persons and 4 households;
  • float, int, date, Enum and str inputs, monthly and ETERNITY;
  • seven calculations over two months.

• Dumps without counts. With count.npy removed, the round trip still holds when the ids give the counts (one per instance, as SimulationBuilder sets them) or the last group has members. Otherwise only the count file gives the count. That case is excluded from the property by assume and tested as intended in test_dump_keeps_a_count_its_ids_and_memberships_do_not_give.

• Counts. build_default_simulation(count=c) raises TypeError exactly when range(c) does, and then leaves the simulation it was given untouched. Otherwise, every population's ids and every group's memberships are list(range(c)) as int64, and population.count == c. test_default_simulation_count_property.py samples it with 300 examples over:

  • Python, NumPy and 0-d integers, signed and unsigned, including negative Python integers;
  • bools and NumPy bools;
  • Python and NumPy floats, 0-d float arrays, NaN and ±inf;
  • strings, None, Decimal and Fraction.

Tests

Run one file at a time: CPython 3.14.7, NumPy 2.4.2 (uv.lock), Hypothesis 6.168.3, macOS arm64, 4 BLAS/OpenMP threads.

File	Upstream master 7399c4ed8	This branch
tests/core/test_dump_restore.py: 11 new tests	11 failed, 1 passed. 4 fail on the restore bugs (no persons; empty groups last; a count only the file gives; no group entities). 6 use COUNT_FILE, which master lacks. 1 (the __index__ count) fails at build, because master's np.arange refuses it.	12 passed
tests/core/test_dump_restore_property.py (new, 200 examples)	collection error (imports COUNT_FILE); with master's restore logic swapped into this branch, it fails	1 passed
tests/core/test_default_simulation_count_property.py (new, 300 examples)	1 failed (minimal case 0.0)	1 passed
tests/core/test_simulation_builder.py: 10 new non-integer cases	10 failed (8 built; "2" and None raised a different TypeError, from np.arange)	45 passed
tests/core/test_group_population_empty.py: the integer-ID test widened from 4 count types to 14 (every NumPy integer width, signed and unsigned, and 0-d integer arrays)	46 passed	46 passed (regression guard)

Mutation checks. I applied each of these to the branch and ran all five files. Each fails at least one test:

• master's restore logic;
• the inferred count without its ids bound;
• the inferred count without its membership bounds;
• no count file written;
• count file ignored;
• the builder without operator.index;
• int32 ids.

The count property's floats are bounded, so a regression fails in about 1 s instead of allocating np.arange of a huge float. An unbounded first version did that at 12:28Z during a mutation check and pushed the host into swap.

Full suite at 619c577: every test file under tests/, one at a time. 2,688 passed, 3 xfailed, 1 skipped, 0 failed. The two tests/smoke files skip locally without the country packages. In CI, the country-package jobs run test_country_init.py; test_us.py's microsimulation is skipped there too (RUN_SMOKE_TESTS: "0" in pr.yaml). After the review fix, the seven files that dump or restore pass at 44e6914:

• test_dump_restore.py: 12;
• test_dump_restore_property.py: 1;
• test_restore_input_registry_property.py: 1;
• test_restore_input_registry.py: 14;
• test_restore_input_provenance.py: 2;
• test_subsample_restore_inputs.py: 8;
• test_option_result_cache.py: 24.

ruff check and ruff format --check pass.

Callers. The local clones of policyengine-us (2a81a31908), policyengine-uk (d292ae10b) and policyengine.py (dd82d0e59) never call dump_simulation, restore_simulation or build_default_simulation directly. The independent review found no direct callers in PE-US either.

Not run locally. The CI matrix (Python 3.11–3.14 on Ubuntu and Windows) and the country-package jobs; CI runs them.

Review

• r1 (GPT-6.1 Sol, at 619c577; ~/reviews/us-hub/reviews-gen/core-624/review-r1.md): REQUEST_CHANGES.
  • P2. The dump saved population.count as given. A masked-array count raised NotImplementedError on dump. An __index__-only object wrote an object array, which restore refuses. Fixed in 44e6914: the dump writes np.int64(operator.index(count)), and writes no file for a count that is not an integer.
  • New regression tests: test_dump_writes_any_integer_count_as_int64 (uint8, 0-d array, masked array, __index__) and test_dump_of_a_count_that_is_not_an_integer_infers_it. Four of these fail at 619c577 and pass now. The 0-d array case already wrote int64 there.
  • Notes answered here:
    • the round-trip invariant's domain is now stated, including the object-dtype ids limit;
    • the property now also compares array type, shape and enum class;
    • the master column now separates the tests that fail on the bugs from those that need COUNT_FILE.
• r2 (GPT-6.1 Sol, delta review at 44e6914; ~/reviews/us-hub/reviews-gen/core-624/review-r2.md): APPROVE, no P1 or P2.
  • The r1 P2 is fixed. Its witnesses round-trip at this head, and a 28-case count matrix passes: integer-like counts write int64 count files, and counts that are not integers write none.
  • The new tests fail at 619c577 (4 failed, 1 passed) and pass here. Mutations of array type, shape, enum class and count-file dtype each fail the property.
  • Its three P3 corrections to this body are applied:
    • the round-trip domain now sets non-integer counts apart;
    • the master failures split as 4 bugs, 6 COUNT_FILE and 1 build;
    • the __index__ wording and the smoke-test coverage in CI are corrected.

Unchanged, noted while testing

• scripts/simulation_generator.make_simulation does not run on master: it builds Simulation without populations, and it reads is_person on a population rather than its entity. The new test builds that structure (no group ids) by hand.
• Object-dtype ids (for example np.array(["h0", "h1"], dtype=object)) can't be restored, at master and here: ValueError: Object arrays cannot be loaded when allow_pickle=False. The builders set numeric or fixed-width string ids.
• An __index__-only count builds again (as before Reduce and rank over a population with no persons #590; master's np.arange refuses it), but set_input then refuses every array, because holders compare an array's length with the count.
• A negative count still builds, with empty ids, as range(-1) did.
• An empty plain list given for an Enum input is float, and is refused with Unsupported array dtype: float64. An empty string or integer array works.

axiom: n/a: core infrastructure (dump/restore and default simulation building), no policy change.

🤖 Generated with Claude Code

[MaxGhenis]

Hand-off to the US + core hub from session local_842b069b-c3ed-4f4f-8583-dcd91e0c85e5, which was started by the hub.

• Head: 44e6914f56fb57e1b198e0cd2f60f11a51b047dc. Not a draft; mergeable is MERGEABLE.
• Independent review: GPT-6.1 Sol.
  • r1 at 619c577 requested changes (one P2: counts dumped as given). Fixed in 44e6914.
  • r2 at 44e6914: APPROVE, no P1 or P2. Its three P3 corrections to the body are applied.
  • Reviews, briefs and evidence: ~/reviews/us-hub/reviews-gen/core-624/ (review-r1.md, review-r2.md, evidence/).
• Passed locally (CPython 3.14.7, NumPy 2.4.2 from uv.lock, one file at a time):
  • full suite at 619c577: 2,688 passed, 3 xfailed, 1 skipped, 0 failed;
  • at 44e6914, the seven files that dump or restore: all pass;
  • fail-before at master and at 619c577, and the mutation checks: see the PR body.
• Pending: CI on 44e6914 (4 checks passed, 15 pending when handed off).
• Depends on: nothing. Reduce and rank over a population with no persons #590 is merged and in the base.
• Holds and decisions: none. This is an engine bug fix with no published numbers. No caller in PE-US, PE-UK or policyengine.py uses the changed functions directly.
• Follow-ups owed: none. The PR body's "Unchanged, noted while testing" lists behaviour that predates this PR.
• Merge gates once CI is green: gh pr checks exits 0 on the latest run, mergeable is MERGEABLE, and there is no CHANGES_REQUESTED review. The repo's merge method applies.

🤖 Generated with Claude Code