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.

Structural reads and scale values (second review)

Formulas sometimes read a node's structure — for state in p.waived_counties._children, set(p.zips._children.keys()), next(iter(p.limits._children.values())), len(node._children). The wrapper now hands out a TracingChildren mapping for _children (and iterates / answers in itself the same way): structure reads are recorded at the node's canonical path with the child names as the value, and children come back through the parent's tracing so nested nodes stay wrapped. No ._children path ever reaches a trace or the map.

Scale reads record {"type", "thresholds", "rates" | "amounts"} as plain lists instead of None, so a serialized trace keeps the schedule.

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.

[anth-volk]

I re-reviewed b916dc51 and confirmed that the four findings from my previous review are addressed. I also tested the proposed changes combined with the current master: 718 tests passed, with one skipped and one expected failure.

Medium: parameter-node structure reads emit invalid parameter paths and can bypass tracing

TracingParameterNodeAtInstant.get_traced_child treats every object that is neither a parameter node nor an allowed scalar, list, or NumPy array as though it were a tax schedule. It constructs a path from the accessed attribute and records None as its value.

PolicyEngine US has four formulas with five reads of the internal _children dictionary:

• is_in_snap_abawd_waived_area
• ut_ccap_income_group (two reads)
• co_ccap_fpg_eligible
• slcsp_rating_area_la_county

I reproduced the resulting dependency-map entry:

"gov.usda.snap.work_requirements.abawd.waived_counties._children": [
  "is_in_snap_abawd_waived_area"
]

_children is an implementation attribute, not a resolvable parameter path. Consequently, the map can expose a key that downstream code cannot resolve as a parameter. There is also a deeper loss in ut_ccap_income_group: p.income_limits._children.values() returns an unwrapped ParameterNodeAtInstant; the later size_node._children read therefore happens outside the tracing wrapper and is not recorded.

Please provide tracing-aware parameter-node iteration or mapping access, or otherwise special-case structural reads so that:

1. the trace records a canonical parameter-node path rather than a path ending in ._children;
2. child names used by the formula are represented;
3. parameter nodes obtained from the mapping remain wrapped; and
4. the dependency map contains only documented, resolvable parameter paths.

A regression test should cover iteration over child names and retrieving a nested node through the child mapping.

Representation note: US tax schedules retain only their path

I also inventoried the current PolicyEngine US parameter tree. Scalar numbers, booleans, lists, and vectorized NumPy selections retain their values. The tree contains 651 SingleAmountTaxScale, 212 MarginalRateTaxScale, and 2 MarginalAmountTaxScale objects; all are recorded as None, so their concrete class, thresholds, rates, and amounts are absent from serialized traces.

That is sufficient for the dependency map because it consumes only parameter paths, but it is not a complete value representation for FullTracer. If complete serialized parameter traces are part of the intended contract, the general tracer should retain a JSON-safe schedule description while _QuietFullTracer can continue discarding values for dependency-map generation.

[PavelMakarchuk]

Thanks for the second pass. Both points are addressed in f140bc2, verified against policyengine-us 1.823.0.

Structural reads (medium). _children is no longer treated like a scale. The wrapper now returns a TracingChildren mapping view for it, and iterating the wrapper itself or testing membership behaves the same way:

1. iterating, len(), keys(), or in records a read at the node's canonical path (gov.usda.snap.work_requirements.abawd.waived_counties, never …._children) with the child names as the value;
2. the names a formula uses are therefore in the trace;
3. [name], .values(), and .items() go through the parent's get_traced_child, so nested nodes stay wrapped and their later reads are recorded;
4. the map contains only resolvable paths.

Regression tests cover iterating child names, retrieving a nested node through the mapping and reading under it, and iterating the node directly. On the four formulas you listed, the map now records gov.usda.snap.work_requirements.abawd.waived_counties plus its state children, gov.states.ut.dwf.ccap.copay.income_limits and the nested household-size node (the read that was previously lost), gov.states.co.ccap.entry.fpg_rate, and gov.aca.la_county_rating_area; a scan of the resulting map finds no key containing ._children.

Scale representation. Taken as part of the contract. A scale read now records {"type", "thresholds", "rates" | "amounts"} as plain lists, so get_serialized_flat_trace() carries the schedule; json.dumps of the serialized trace is exercised in the test. _QuietFullTracer still discards calculation results; parameter values are small and kept.

Full suite: 684 passed, 4 skipped, 1 xfailed, ruff clean.

[PavelMakarchuk]

@anth-volk re-requesting your review on f140bc2, which addresses the second round (structural _children reads and scale values in traces); CI is green.

One question that could shorten this: in your first review you suggested splitting the tracer corrections (#541, #542, plus the branch and structural-read fixes) from the dependency-map command so the narrower fix can land on its own. Do you still want that? If so I'll split the branch and open the tracer-only PR first.

[anth-volk]

Let's keep this as one PR. My earlier note about splitting it was only an optional suggestion, not a required change. I'll review the updated implementation together at f140bc2.