Trace parameter reads in yearly formulas and scale reads; add dependency-map

[PavelMakarchuk]

Fixes #541. Fixes #542.

The two tracer gaps

With simulation.trace = True, TraceNode.parameters was almost always empty for real models:

• FullTracer records no parameter reads for yearly formulas: the at-instant node is cached before the tracing recast #541 — _run_formula recast the parameter tree for tracing lazily, but by then defined_for and adds/subtracts evaluation had already built and cached the yearly ParameterNodeAtInstant as a plain node. Every parameters(period).gov.* read for that instant bypassed TracingParameterNodeAtInstant. Monthly instants were created later and traced fine, which hid the bug.
• TracingParameterNodeAtInstant never records scale/bracket parameter reads #542 — TracingParameterNodeAtInstant.get_traced_child only recorded scalar and array leaves, so scale and bracket reads (p.base.calc(age), p.max[children], every bracket schedule) were never recorded.

The fix

• ParameterNode.set_tracing(tracer, branch_name) recasts the tree and clears every at-instant cache in the subtree whenever tracing switches on or off (clear_at_instant_caches).
• The Simulation.trace setter calls it (and is now the only place self.tracer is built; the duplicate assignment in __init__ is gone). _run_formula only re-points the tree at the current branch, since branch simulations share the system. Systems with no parameter tree are skipped.
• Non-node, non-leaf children (scales, brackets) are recorded at their node with value=None, matching how vectorial reads already record at the parent.

Three regression tests in test_tracers.py fail on master and pass here: the yearly read after the node was cached untraced, the scale read, and trace off restoring plain nodes.

policyengine-core dependency-map

Builds on the fixed tracer. policyengine-core dependency-map -c policyengine_us traces which variables read each parameter and which variables feed which, and writes them as readers and consumers with the package version, a sha256 fingerprint over its model surface (entities, parameters, system, variables), and the core version. Downstream: validation matching in policyengine-app-v2 (PolicyEngine/policyengine-app-v2#1180 consumes this shape), the calibration dashboard's model-coverage page (which derives the same edges by grepping source today), reform classifiers.

Populations: tests (default) traces the package's YAML tests, one per newly covered output variable per file; --every-test traces all; microdata traces every variable over a subsample of the package's default dataset; both unions them. Reads of gov.abolitions.* are skipped: core's neutralisation check reads them before the formula, so they are a switch per variable, not a dependency.

Country packages need no wrapper. The release job that publishes the map per version belongs in each country repo and is a follow-up.

Design note: no cache clearing

The first version of this branch cleared the parameter at-instant caches whenever tracing switched or the tracer changed. Profiling against policyengine-us showed each clear rebuilds the whole parameter tree (~2 s), and branch simulations swap tracers per formula, so a traced microsim run went from 11 s to 29 s and a traced household calculation from 0.3 s to 4 s. The cache now holds plain nodes only; _get_at_instant wraps one in TracingParameterNodeAtInstant on the way out when tracing is on, and the wrapper reads the tracer and branch name from the parameter root at access time. Nothing is ever invalidated, set_tracing is three attribute writes, and parameter reads are now labelled with the current branch rather than whichever built the cache.

Benchmark (policyengine-us 1.822.4, 2,000-household subsample, 19 output variables; household = one CA family, household_net_income + income_tax)

	released 3.32.1	this branch
microsim, trace off	10.7 s	10.0 s
microsim, trace on	11.0 s	11.1 s
parameter reads recorded	0	75,874
household, trace off	0.24 s	0.28 s
household, trace on	0.28 s	0.35 s
household parameter reads in flat trace	0	8,154

Same trace node count (120,947) and flat-trace entries (6,993) on both, so the tracer's variable-level output is unchanged; the only difference is that parameter reads are now present. get_flat_trace, get_serialized_flat_trace, and computation_log.lines() all run on the traced household.

Downstream consumers

policyengine-api (its own pins: policyengine-us 1.764.6), running PolicyEngineCountry.calculate, which sets trace = True and returns computation_log.lines(aggregate=False, max_depth=10), on a two-person California household requesting income tax, refundable CTC, EITC, SNAP, SPM net income, and household net income:

	core 3.30.1 (current pin)	core 3.32.1	this branch
computed values	baseline	identical	identical
computation log	5,231 lines	byte-identical	byte-identical
3 calls, seconds	4.9 / 1.3 / 5.5	2.0 / 0.5 / 2.2	2.1 / 0.6 / 2.1

The API's household-calculation unit tests (test_household_calculation_service, test_calculate_*, test_household_traces: 17 tests) pass on this branch.

policyengine-us 1.822.4 on this branch: the Python suite (tests/ minus policy/ and microsimulation/) passes except five tests in test_run_selective_tests.py, which fail identically on released core (they shell out to git and were run from site-packages). Its YAML tests for gov/irs, gov/states/ca, and gov/usda through this branch's policyengine-core test runner: 2,504 passed in 6m46s. Itemisation branching (simulation.get_branch) is exercised by every income-tax calculation above.

Verified

• Full core suite: 676 passed, 4 skipped. The four test_release_tagging.py failures are pre-existing on master (same result with this diff stashed).
• Country-template YAML tests via the CLI: 39 passed.
• Against policyengine-us 1.822.4 in a scratch venv with this core installed, policyengine-core dependency-map -c policyengine_us --tests-root <ctc tests> records gov.irs.credits.ctc.amount.base → ctc_child_individual_maximum (scale read) and gov.irs.credits.ctc.refundable.fully_refundable → refundable_ctc (yearly read), with no monkeypatches.

🤖 Generated with Claude Code

[anth-volk]

I am requesting changes because the direct fixes for #541 and #542 work, but the new dependency map still omits real reform dependencies and silently skips a supported YAML output form. I also reproduced incorrect branch metadata in serialized parameter traces.

Findings

1. High: parameter-backed adds and subtracts are absent from traces

When a variable uses a parameter path for adds or subtracts, _run_formula resolves the Parameter directly and calls it with period.start (simulation.py lines 1052–1064). That call bypasses TracingParameterNodeAtInstant.

I reproduced this with an actual US reform:

• Reform: set gov.irs.credits.refundable to an empty list for 2026.
• Baseline income_tax_refundable_credits: 5065.6220703125.
• Reformed value: 0.0.
• Parameters recorded on the reformed income_tax_refundable_credits trace node: only gov.abolitions.income_tax_refundable_credits; gov.irs.credits.refundable is absent.

The same direct-access problem applies to parameter-backed subtracts, and variable-level uprating also resolves parameters outside the tracing wrapper. These are material parameter dependencies: changing the list changes which variables contribute to the result. A map intended to classify reforms needs to record them explicitly in the active trace node or document and compensate for their omission.

2. High: entity-scoped YAML outputs are silently skipped

trace_yaml_tests loops over the top-level keys of test["output"] and passes each key to simulation.calculate (dependency_map.py lines 136–150). The standard YAML runner also supports entity-scoped output structures such as:

output:
  people:
    person1:
      al_ui: 3850

For these cases the new command attempts to calculate a variable named people, suppresses the exception, and never calculates the nested variables. The deduplication logic remembers only the top-level people key, so it also treats later cases in the same file as already covered.

In the current policyengine-us repository, 206 of the 5,405 cases selected by the default dependency-map logic have entity-scoped outputs. Running the command on policyengine_us/tests/policy/baseline/gov/states/al/dol/unemployment_insurance/al_ui.yaml reported:

tests: 1
failed: 0
parameter paths: 0
consumed variables: 0

The file contains valid nested person outputs. Output-variable extraction should follow the same entity-aware and period-aware rules as YamlItem.check_output. Calculation exceptions should also increment failure statistics or be surfaced, otherwise a materially incomplete map appears successful.

3. Medium: a nested branch can relabel later parent parameter reads

TracingParameterNodeAtInstant.tracer and .branch_name read mutable state from the shared parameter root at access time (tracing_parameter_node_at_instant.py lines 41–51). A nested branch calculation changes that shared state in _run_formula, but the parent state is not restored when the nested calculation returns.

I reproduced this through the existing ctc_limiting_tax_liability formula. Its trace node had branch default, while its later read of gov.irs.credits.non_refundable was labeled no_salt, the branch that had just completed. The dependency map currently ignores branch names, but the serialized tracer output is incorrect.

Now that _at_instant_cache contains only plain nodes, a wrapper returned for a particular formula can capture that formula's tracer and branch instead of consulting mutable root state. Alternatively, the active tracing state needs scoped restoration after nested calculations.

4. Low: the release note and source comment describe removed cache clearing

The final implementation deliberately keeps plain at-instant nodes cached, but the changelog fragment still says the caches are cleared (fragment). The comment in the Simulation.trace setter says “Recast (and clear)” as well. These should describe the final wrap-on-access design.

Successful verification

I tested PR head de2b4090533608c640900e31596199ac6de9198a against current policyengine-us 1.823.0 on Python 3.14.6.

• A reform setting gov.irs.credits.ctc.amount.base[0].amount to $2,500 produced [0, 2500] for ctc_child_individual_maximum. get_serialized_flat_trace() recorded gov.irs.credits.ctc.amount.base with null, the intended representation for a scale.
• A reform setting gov.irs.credits.ctc.refundable.fully_refundable to true produced a $2,500 refundable CTC and recorded that yearly boolean parameter with value true.
• A vectorized federal-poverty-guideline calculation recorded gov.hhs.fpg.first_person and gov.hhs.fpg.additional_person with their selected array values.
• A Social Security PIA calculation recorded the marginal-rate scale gov.ssa.social_security.pia.formula_factors with null.
• A CTC-only dependency-map run completed over 28 selected YAML cases, producing 319 parameter paths and 877 consumed variables. It included both gov.irs.credits.ctc.amount.base and gov.irs.credits.ctc.refundable.fully_refundable.

Commands run:

uv run --frozen pytest -q tests/core/test_parameters.py tests/core/parameters_fancy_indexing tests/core/test_reforms.py tests/core/test_reform_parameter_isolation.py tests/core/test_reform_period_keys.py tests/core/test_simulations.py tests/core/test_simulation_builder.py tests/core/test_tracers.py tests/core/tools/test_dependency_map.py
# 170 passed

uv run --frozen --extra dev ruff check <changed Python files and tests>
# passed

uv run --frozen --extra dev ruff format --check <changed Python files and tests>
# 7 files already formatted

policyengine-core dependency-map -c policyengine_us --tests-root policyengine_us/tests/policy/baseline/gov/irs/credits/ctc --output /tmp/ctc-dependency-map.json
# 28 tests, 0 reported failures, 319 parameter paths, 877 consumed variables

The GitHub checks are also currently successful across the reported Python and operating-system matrix. I did not run the full US microdata map because the deterministic omissions above already establish that its output can be incomplete.

Overall assessment and documentation review

Keeping plain at-instant nodes cached and wrapping them on access is a good direction: it fixes the yearly-read timing problem without the measured cache-rebuild cost. Recording tax-scale objects at their parameter path with a null value also matches the existing vectorized-access granularity.

The tracer changes are close, but the branch context needs to remain associated with the formula that obtained the wrapper. The dependency-map command should not be treated as complete until it handles core-managed parameter dependencies and the YAML runner's supported output structures. Splitting the tracer correction from the dependency-map command would also let the narrower #541/#542 work proceed independently.

Documentation observed: module documentation, command-line help, and two changelog fragments. Additional documentation is needed for the output schema and for how partial calculation failures affect completeness, and the stale cache-clearing statements need correction. Impact: high. Confidence: high, based on direct US reform and command execution. Known gap: I did not run the full microdata population after reproducing the deterministic correctness problems above.

[PavelMakarchuk]

Thanks for the review. All four findings reproduced against the code and are addressed in b916dc5.

1. Parameter-backed adds/subtracts and uprating (high). Confirmed: _run_formula resolves those on the tree directly and variable.uprating does the same in _calculate, so none of them passed through TracingParameterNodeAtInstant. They're now recorded on the variable's trace node via Simulation._record_parameter_read: the list parameter itself (with the list as its value), each parameter path inside an adds/subtracts list (with its value), and both instants an uprating factor reads. Tests: test_parameter_backed_adds_and_subtracts_are_traced, test_uprating_parameter_reads_are_traced.

2. Entity-scoped YAML outputs (high). Confirmed. iter_output_variables now reads output the way YamlItem.check_output does: a key is a variable, an entity's singular key holding variables, or an entity's plural key holding instances of variables, and a dict value keyed by period checks the variable at each period. Deduplication uses the flattened variable names, so a later case is no longer hidden by a shared people key. Output calculations that raise are counted as outputErrors (with outputs alongside), the command prints both, and the module docstring now documents the output schema and states that failed/outputErrors bound how incomplete a map may be. Tests: test_trace_yaml_tests_follows_entity_scoped_and_period_keyed_outputs, test_trace_yaml_tests_counts_outputs_that_fail_to_calculate.

3. Nested branch relabelling (medium). Confirmed via your ctc_limiting_tax_liability scenario reproduced on the template. The wrapper now captures its tracer and branch when created (the root-state indirection is gone, as you suggested, since nothing caches wrappers any more), _run_formula restores the tree's tracing state after the formula returns, and get_branch hands the tree back to the caller after the branch's trace setter re-pointed it. Test: test_nested_branch_does_not_relabel_the_callers_parameter_reads asserts the caller's reads before and after the nested calculation are labelled default, the branch's read is labelled inner, and the tree is back on the caller afterwards. test_swapping_tracer_keeps_cached_nodes_and_records_into_new_tracer was rewritten for the capture-at-creation semantics.

4. Stale cache-clearing wording (low). Changelog fragment and the Simulation.trace setter comment now describe the wrap-on-access design.

Checks: full core suite 711 passed / 4 skipped / 1 xfailed; your focused set 175 passed; ruff check and format clean on the changed files; policyengine-core test on the template 39 passed; policyengine-core dependency-map -c policyengine_core.country_template now prints tests: 10, failed: 0, outputs: 15, outputErrors: 0.

On splitting the tracer fix from the command: happy to do that if you'd prefer it for the #541/#542 release; leaving it as one PR for now so the command's tests keep exercising the tracer changes.