Keep the record of set_input values in step with storage

[MaxGhenis]

Fixes #559.

Problem

Simulation._user_input_keys records each (variable, branch, period) a simulation stored through set_input. In core, two places read it: _invalidate_all_caches (run by apply_reform) keeps the values it names, and to_input_dataframe exports them. In 3.32.12 the record drifted from storage:

1. delete_arrays deleted values but kept their entries. A formula result calculated later for the same period then counted as an input. It survived apply_reform and was exported by to_input_dataframe; an input variable read after deletion exported its default. This happened through Simulation.delete_arrays and through Holder.delete_arrays, which country packages call directly. policyengine-us system.py, for example, moves employment_income into employment_income_before_lsr that way, and on master the record kept ('employment_income', 'default', 2025) after the move.
   • With auto_carry_over_input_variables, the surviving formula result is also carried into later periods. The review of Carry over only inputs, the latest at or before the requested period #562 ran this case: a variable whose formula gives 7 through 2012 is set to 20 for 2012, deleted, and calculated (7). After apply_reform, master gives 7 for 2013, in memory and on disk. A simulation that never had the input gives 0. test_formula_result_for_a_deleted_input_is_not_carried_over pins it for both storages.
2. clone() (so also get_branch) shared one record between simulations that store their values separately. As a result:
   • an input set on a clone was recorded for the original;
   • an input set on a parent after a branch was made was recorded for the branch, whose storage never received it;
   • the list of running set_input calls was shared too.
3. Values a custom set_input handler calculated were recorded as inputs. A handler that calls calculate before storing its input made those formula results survive apply_reform.
4. Entries named the period as given, not as stored.
   • An eternal input set for 2025 was recorded as 2025 while storage keys it as eternity, so deleting it left its entry.
   • A handler storing a month as the string "2025-01" was recorded with a string, so to_input_dataframe did not export it.
   • Twelve months starting on the first of a month (month:2025-01:12, month:2025-03:12) were recorded as twelve months, while storage keys them as the year starting then (2025, year:2025-03). Deleting the value left its entry.
5. subsample rebuilt every stored value but kept the old record.

Change

• Recording. Holder._set records the period storage keys the value under. Storage keys a value by its period's string form, so the recorded period is that string read back: eternity for an eternal variable, the year for twelve months starting on the first of a month, and otherwise the period itself. Each entry therefore names one stored value. _set takes is_input: Optional[bool] = None, the same keyword and meaning as Make set_input on a branch drop values calculated from the input it replaces #560: None means "an input if a set_input call is running". put_in_cache passes False, so calculated values are never recorded and never redirected to the input's branch.
• Deletion. Holder.delete_arrays notes the keys memory and disk storage hold before deleting, and afterwards finds the ones that are gone. It drops the entry for a removed key only if neither storage still holds a value for it. Simulation.delete_arrays calls it for each visible branch.
  • The cost depends on what the holder stores, not on the size of the record: the record is never iterated, with or without disk storage.
  • Whether a storage still holds a value is read from its keys, so no stored file is loaded.
  • Disk storage deletes only the exact period it is given (OnDiskStorage.delete(period) deletes only that exact period, not the periods within it #564), so a value deleted from memory can survive on disk. Its entry stays, as on master.
• Copies. Simulation.clone gives the copy set(self._user_input_keys) and an empty list of running set_input calls.
• subsample. It starts the record again before it rebuilds the simulation.

Unchanged: storage, what set_input stores, and the values apply_reform keeps for inputs that were not deleted (test_inputs_not_deleted_are_still_kept_by_apply_reform, test_input_that_was_not_deleted_is_carried_over_after_apply_reform). calculate does not read the record. A simulation that never calls apply_reform / _invalidate_all_caches after deleting, cloning or running a calculating handler, and never exports, therefore calculates the same values.

• policyengine-uk main and policyengine.py main do not read _user_input_keys (git grep). policyengine-us main reads it in one place: Simulation._rebind_holders (policyengine_us/spm.py) replaces a simulation's record with a new set holding the entries for variables its system has. It was written when clones shared one record; with this PR the clone already has its own, and the filter works the same. PE-US's two tests of it pass on this head (see Tests and checks).
• policyengine.py main never calls apply_reform, to_input_dataframe, subsample or delete_arrays (git grep).
• In their constructors, PE-US (system.py) and PE-UK (simulation.py) call apply_reform for structural reforms before they move inputs such as employment_income with delete_arrays. An apply_reform after such a move and a calculation now recalculates the moved variable; on master it kept the pre-reform formula result.

One intended change for handler authors: a value a custom set_input handler stores through put_in_cache is a cached calculation, not an input. Handlers store inputs with holder._set or holder.set_input, as core's own set_input_divide_by_period and set_input_dispatch_by_period do. No handler in policyengine-us, policyengine-uk or policyengine-canada calls put_in_cache (git grep).

Performance

delete_arrays on every variable of a branch, as country marginal-rate code does, with a 9,000-entry record (3,000 input variables × 3 years) and 3,024 variables. CPU seconds, median of 5 loops (takeover/record_size_bench.out):

Core	Memory only	Disk storage configured, empty
master (no pruning)	0.006	0.008
a9bdeff (this PR before the round-2 fixes)	0.020	1.097
this PR	0.019	0.022

test_delete_does_not_look_through_the_whole_record and test_disk_delete_reads_no_file_and_does_not_look_through_the_record make the record raise on iteration and make disk storage raise on any file read, then delete.

Invariants

These hold for every sequence of operations on a family of simulations (a root, its clones and branches):

1. The record follows storage. Every entry names a value the simulation stores, under the entry's branch and period. In memory it equals what set_input stored.
2. The record matches a reference model. Each simulation's entries are exactly the values it stored through set_input that some storage still holds, including those it inherited when it was cloned or branched. Values calculated during a set_input call are not entries.
3. Isolation. An operation on one simulation never changes another simulation's record.
4. _invalidate_all_caches keeps exactly the inputs. Afterwards, the simulation and its branches store their inputs and nothing else.
5. Export equals inputs. For an input variable, to_input_dataframe exports the recorded periods of the variable's own unit on the branches the simulation reads, with the input's values, and stores nothing.

tests/core/test_user_input_keys_property.py has two properties:

• test_user_input_keys_match_reference_model checks 1, 2, 3 and 5 after every step, and 4 after every _invalidate_all_caches, on memory storage. It runs 300 random sequences of up to 30 steps.
  • The steps are set_input, calculate, Simulation.delete_arrays, Holder.delete_arrays, clone, get_branch and _invalidate_all_caches.
  • The variables are monthly person and household variables, an eternal variable, a monthly variable set for twelve months (calendar and rolling year), and a variable whose yearly set_input handler calculates income_tax and then stores three months under string periods.
  • The model is seeded from the situation, not from the implementation's record, and the export check reads to_input_dataframe itself.
• test_user_input_keys_follow_memory_and_disk_storage checks 1, 2 and 4 on one simulation whose holders store in memory or on disk, switching between the two at random. It runs 200 sequences of up to 25 steps. Its model takes "some storage still holds the value" from the storages' keys, so it holds however disk storage deletes.
• On master both fail at the first step with set_input('birth', '2025'). On a9bdeff both fail with set_input('salary', 'month:2025-01:12').
• The module skips itself without Hypothesis, as the smoke job requires.

Mutation check at this head (takeover/mutants.py): 16 mutants, each killed by at least one of the two modules. They are:

• no pruning call;
• taking the "before" keys after deleting;
• no memory comparison;
• no disk comparison;
• splitting disk keys on the first _;
• dropping a removed key's entry without checking what storage still holds;
• that check looking at memory only, or at disk only;
• that check loading the file;
• recording the raw period;
• no string round trip for twelve-month periods;
• no eternity in _storage_period;
• put_in_cache recording inputs;
• subsample keeping the record;
• no record copy on clone;
• no reset of running set_input calls.

Tests and checks

• tests/core/test_user_input_keys.py: 29 example tests. On master 25 fail; the 4 that pass guard against dropping too much. On a9bdeff the 5 for the round-2 findings fail. Both properties fail on master and on a9bdeff.
• CI at 579c375: 18/18 checks pass. pytest tests: 1,174 passed, 4 skipped, 1 xfailed; country-template YAML: 39 passed (Ubuntu, Python 3.13).
• Smoke-job collection without Hypothesis: the property module skips itself (the CI "Test Core and country packages" jobs pass).
• Compositions at 579c375: each open core PR merged with this head, then this PR's two modules plus that PR's own test modules (takeover/compose_focused.out):
  • Read only periods the current branch can see when uprating or carrying over #552 (5fe8b61): 563 passed;
  • Carry over only values stored at the variable's definition period #557 (97d4582): 275 passed;
  • Give each disk-backed holder storage its own directory #558 (a9a0b40): 69 passed;
  • Delete the periods within a deleted period from disk storage #565 (bfa78cb, disk storage deletes contained periods): 623 passed;
  • Carry over only inputs, the latest at or before the requested period #562 (912a4ac): 88 passed, after resolving the one conflict (Holder._set / put_in_cache, both sides kept);
  • Uprate only from inputs, so uprated values don't depend on calculation order #563 (2120f78, stacked on Carry over only inputs, the latest at or before the requested period #562): 355 passed, same resolution;
  • Make set_input on a branch drop values calculated from the input it replaces #560 (116b018, draft) conflicts in Holder._set / put_in_cache / delete_arrays and Simulation.clone / subsample, with both sides kept; its owner takes the resolution when rebasing.
• The carry-over case on Keep the record of set_input values in step with storage #561 + Carry over only inputs, the latest at or before the requested period #562 combined: 0 for 2013 in memory and on disk (takeover/repro_compose_562.out). On master: 7 in both.
• Single-year microsimulations, master b78b0ba against 579c375: running (Enhanced FRS 2026 on policyengine-uk main; a 3,000-household Enhanced CPS 2026 sample with marginal tax rates on policyengine-us main). Results will be added here before merge.

make documentation was not run locally; CI's Test jobs build the documentation, and they pass.

Review

• Round 1 (GPT-6.1 Sol, on 390e9e8): no regression introduced; six failures that already exist on master. Three were fixed in a9bdeff (calculating handlers, string periods, subsample); three are disk-storage bugs listed below.
• Round 2 (on a9bdeff): the Opus code review requested changes and the Sol soundness review said unsound. All findings are answered in this head:
  • a value held in memory and on disk lost its entry when only memory's copy was deleted: fixed, test_entry_is_kept_while_disk_still_holds_a_value_deleted_from_memory;
  • twelve-month periods were recorded differently from how storage keys them: fixed, test_twelve_month_input_is_recorded_as_the_year_it_is_stored_as;
  • a holder with disk storage configured went through the whole record and loaded surviving files on every delete: fixed, see Performance;
  • the property model was seeded from the implementation and did not read to_input_dataframe: fixed;
  • a handler that stores its input through put_in_cache is no longer recorded: intended, see Change;
  • calculate_add over a period that holds an input (a monthly variable given a two-month input) overwrites the input with the sum, on master and here: storage behaviour, not the record. Carry over only inputs, the latest at or before the requested period #562's put_in_cache guard keeps the input.
• Round 3 (on this head): in progress (an Opus code review and a GPT-6.1 Sol soundness review).

Not in this PR

Pre-existing disk-storage bugs that change storage itself, not the record:

• OnDiskStorage.clone keeps the same .npy paths, so a clone writing a key overwrites the original's file (Give each disk-backed holder storage its own directory #558 gives each holder storage its own directory).
• Deleting a branch name deletes files of every branch whose name starts with it plus _, and get_known_branch_periods splits on every _. Read only periods the current branch can see when uprating or carrying over #552 parses the key on its last _, which fixes both.
• OnDiskStorage.delete(period) deletes only that exact period (OnDiskStorage.delete(period) deletes only that exact period, not the periods within it #564, fixed in Delete the periods within a deleted period from disk storage #565). InMemoryStorage.delete deletes the periods within it, as Holder.delete_arrays documents.

The disk tests here compare the record with what storage still holds, so they pass before and after #552 and #565.

#560 moves _invalidate_all_caches onto per-array input flags. to_input_dataframe still reads _user_input_keys, so this fix is needed either way. _set's is_input keyword is shared, and the two compose. #562 adds a derived keyword next to it; the resolution keeps both (takeover/compose_562_resolution.diff).

axiom: n/a: infrastructure (simulation engine), no policy rule.

🤖 Generated with Claude Code