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

changelog.d/dependency-map.added.md

@@ -0,0 +1 @@
+`policyengine-core dependency-map -c <country package>` traces which variables read each parameter and which variables feed which, over the package's YAML tests and optionally its microdata, and writes the map with the package's version and fingerprint for downstream validation tools.

changelog.d/tracer-parameter-reads.fixed.md

@@ -0,0 +1 @@
+`FullTracer` now records parameter reads in yearly formulas: the parameter tree is pointed at the tracer when `simulation.trace` is set instead of lazily inside `_run_formula`, after `defined_for` and `adds`/`subtracts` evaluation had already read the yearly node (#541). The at-instant cache keeps plain nodes and wraps one for tracing on each access, so switching tracing or swapping branch tracers rebuilds nothing. Scale and bracket reads (`p.rates.calc(x)`, `p.max[children]`) are recorded at their node too (#542). Parameter-backed `adds`/`subtracts` lists and uprating factors, which the core resolves outside the formula, are recorded on the variable's trace node as well, and a nested branch calculation no longer relabels the parameter reads its caller makes after it returns.

policyengine_core/parameters/parameter_node.py

@@ -214,17 +214,19 @@ def clone(self) -> "ParameterNode":
         return clone
 
     def _get_at_instant(self, instant: Instant) -> ParameterNodeAtInstant:
-        if instant in self._at_instant_cache:
-            return self._at_instant_cache[instant]
-        node_at_instant = ParameterNodeAtInstant(self.name, self, instant)
+        # The cache holds plain nodes only; tracing wraps one on the way out.
+        # Building the at-instant tree is expensive, so it must survive
+        # tracing being switched on and off (the API traces some household
+        # simulations and not others on one shared tax-benefit system).
+        node_at_instant = self._at_instant_cache.get(instant)
+        if node_at_instant is None:
+            node_at_instant = ParameterNodeAtInstant(self.name, self, instant)
+            self._at_instant_cache[instant] = node_at_instant
         if self.trace:
-            at_instant = TracingParameterNodeAtInstant(
+            return TracingParameterNodeAtInstant(
                 node_at_instant, self.tracer, self.branch_name
             )
-        else:
-            at_instant = node_at_instant
-        self._at_instant_cache[instant] = at_instant
-        return at_instant
+        return node_at_instant
 
     def attach_to_parent(self, parent: "ParameterNode"):
         self.parent = parent
@@ -234,6 +236,18 @@ def clear_parent_cache(self):
         if self.parent is not None:
             self.parent.clear_parent_cache()
 
+    def set_tracing(self, tracer, branch_name: str) -> None:
+        """Route parameter reads through ``tracer`` (``None`` to stop).
+
+        Cheap by design: the at-instant cache holds plain nodes and
+        ``_get_at_instant`` wraps one with the current ``tracer`` and
+        ``branch_name`` on the way out, so nothing is rebuilt when tracing
+        is switched or a branch simulation swaps tracers.
+        """
+        self.trace = tracer is not None
+        self.tracer = tracer
+        self.branch_name = branch_name
+
     def mark_as_modified(self):
         self.modified = True
         if self.parent is not None:

policyengine_core/scripts/policyengine_command.py

@@ -105,12 +105,43 @@ def build_data_parser(parser):
 
         return parser
 
+    def build_dependency_map_parser(parser):
+        parser = add_tax_benefit_system_arguments(parser, country_only=True)
+        parser.add_argument(
+            "--population",
+            choices=["tests", "microdata", "both"],
+            default="tests",
+            help="what to trace: the package's YAML tests (default), a microdata subsample, or both",
+        )
+        parser.add_argument(
+            "--tests-root",
+            default=None,
+            help="directory of YAML tests to trace (default: <package>/tests)",
+        )
+        parser.add_argument(
+            "--every-test",
+            action="store_true",
+            help="trace every test instead of one per newly covered output variable per file",
+        )
+        parser.add_argument(
+            "--households", type=int, default=2000, help="microdata subsample size"
+        )
+        parser.add_argument("--year", type=int, default=2026, help="microdata year")
+        parser.add_argument("--output", default="dependency-map.json")
+        return parser
+
     parser_test = subparsers.add_parser("test", help="Run OpenFisca YAML tests")
     parser_test = build_test_parser(parser_test)
 
     parser_data = subparsers.add_parser("data", help="Manage OpenFisca data")
     parser_data = build_data_parser(parser_data)
 
+    parser_dependency_map = subparsers.add_parser(
+        "dependency-map",
+        help="Trace which variables read each parameter and which variables feed which",
+    )
+    parser_dependency_map = build_dependency_map_parser(parser_dependency_map)
+
     return parser
 
 
@@ -129,6 +160,11 @@ def main():
 
         return sys.exit(main(parser))
 
+    if args.command == "dependency-map":
+        from policyengine_core.tools.dependency_map import main
+
+        return sys.exit(main(parser))
+
 
 if __name__ == "__main__":
     sys.exit(main())

policyengine_core/simulations/simulation.py

@@ -18,11 +18,7 @@
 from policyengine_core.periods import Period
 from policyengine_core.periods.config import ETERNITY, MONTH, YEAR
 from policyengine_core.periods.helpers import period
-from policyengine_core.tracers import (
-    FullTracer,
-    SimpleTracer,
-    TracingParameterNodeAtInstant,
-)
+from policyengine_core.tracers import FullTracer, SimpleTracer
 import random
 from policyengine_core.tools.hugging_face import *
 from policyengine_core.tools.google_cloud import (
@@ -152,8 +148,7 @@ def __init__(
         # simulation was loaded from.
         self._user_input_keys: set[tuple[str, str, Period]] = set()
         self.debug: bool = False
-        self.trace: bool = trace
-        self.tracer: SimpleTracer = SimpleTracer() if not trace else FullTracer()
+        self.trace: bool = trace  # also builds self.tracer
         self.opt_out_cache: bool = False
         # controls the spirals detection; check for performance impact if > 1
         self.max_spiral_loops: int = 10
@@ -542,6 +537,16 @@ def trace(self, trace: SimpleTracer) -> None:
             self.tracer = FullTracer()
         else:
             self.tracer = SimpleTracer()
+        # Parameter reads are recorded through the parameter tree, which
+        # wraps its cached plain at-instant nodes for tracing on access.
+        # Point it at the tracer here rather than lazily in _run_formula:
+        # by the time a formula runs, defined_for and adds/subtracts
+        # evaluation has already read the yearly node.
+        parameters = getattr(
+            getattr(self, "tax_benefit_system", None), "parameters", None
+        )
+        if parameters is not None:
+            parameters.set_tracing(self.tracer if trace else None, self.branch_name)
 
     def link_to_entities_instances(self) -> None:
         for _key, entity_instance in self.populations.items():
@@ -855,6 +860,14 @@ def _calculate(self, variable_name: str, period: Period = None) -> ArrayLike:
                         )
                     value_in_last_period = uprating_parameter(latest_known_period.start)
                     value_in_this_period = uprating_parameter(period.start)
+                    self._record_parameter_read(
+                        uprating_parameter.name,
+                        latest_known_period.start,
+                        value_in_last_period,
+                    )
+                    self._record_parameter_read(
+                        uprating_parameter.name, period.start, value_in_this_period
+                    )
                     if value_in_last_period == 0:
                         uprating_factor = 1
                     else:
@@ -1034,6 +1047,19 @@ def calculate_output(self, variable_name: str, period: Period = None) -> ArrayLi
 
         return variable.calculate_output(self, variable_name, period)
 
+    def _record_parameter_read(self, name: str, instant, value) -> None:
+        """Record a parameter the core read on a variable's behalf.
+
+        Parameter-backed ``adds``/``subtracts`` lists and uprating factors
+        are resolved on the tree directly rather than through the formula's
+        ``parameters(period)`` wrapper, so they would otherwise be missing
+        from the variable's trace node.
+        """
+        if self.trace:
+            self.tracer.record_parameter_access(
+                name, str(instant), self.branch_name, value
+            )
+
     def _run_formula(
         self, variable: str, population: Population, period: Period
     ) -> ArrayLike:
@@ -1056,6 +1082,9 @@ def _run_formula(
                             f"In the variable '{variable.name}', the 'adds' attribute is a string '{variable.adds}' that does not match any parameter."
                         )
                     adds_list = adds_parameter(period.start)
+                    self._record_parameter_read(
+                        adds_parameter.name, period.start, list(adds_list)
+                    )
                 else:
                     adds_list = variable.adds
                 values = 0
@@ -1070,7 +1099,11 @@ def _run_formula(
                                 self.tax_benefit_system.parameters,
                                 added_variable,
                             )
-                            values = values + parameter(period.start)
+                            added_value = parameter(period.start)
+                            self._record_parameter_read(
+                                parameter.name, period.start, added_value
+                            )
+                            values = values + added_value
                         except:
                             raise ValueError(
                                 f"In the variable '{variable.name}', the 'adds' attribute is a list that contains a string '{added_variable}' that does not match any variable or parameter."
@@ -1087,6 +1120,9 @@ def _run_formula(
                             f"In the variable '{variable.name}', the 'subtracts' attribute is a string '{variable.subtracts}' that does not match any parameter."
                         )
                     subtracts_list = subtracts_parameter(period.start)
+                    self._record_parameter_read(
+                        subtracts_parameter.name, period.start, list(subtracts_list)
+                    )
                 else:
                     subtracts_list = variable.subtracts
                 if values is None:
@@ -1104,29 +1140,38 @@ def _run_formula(
                                 self.tax_benefit_system.parameters,
                                 subtracted_variable,
                             )
-                            values = values - parameter(period.start)
+                            subtracted_value = parameter(period.start)
+                            self._record_parameter_read(
+                                parameter.name, period.start, subtracted_value
+                            )
+                            values = values - subtracted_value
                         except:
                             raise ValueError(
                                 f"In the variable '{variable.name}', the 'subtracts' attribute is a list that contains a string '{subtracted_variable}' that does not match any variable or parameter."
                             )
             return values
 
-        if self.trace and not isinstance(
-            self.tax_benefit_system.parameters, TracingParameterNodeAtInstant
-        ):
-            # Soft-recast
-            self.tax_benefit_system.parameters.branch_name = self.branch_name
-            self.tax_benefit_system.parameters.trace = True
-            self.tax_benefit_system.parameters.tracer = self.tracer
         parameters_at = self.tax_benefit_system.parameters
+        previous_tracing = None
+        if self.trace and parameters_at is not None:
+            # Point the shared parameter tree at this simulation's tracer and
+            # branch for the duration of the formula, then hand it back: a
+            # nested branch calculation must not relabel the reads its
+            # caller makes after it returns.
+            previous_tracing = (parameters_at.tracer, parameters_at.branch_name)
+            parameters_at.set_tracing(self.tracer, self.branch_name)
 
         # A rules-engine formula must be a pure, deterministic function of its
         # inputs. Randomness is forbidden statically at variable registration
         # (check_formula_determinism), so no runtime guard is needed here.
-        if formula.__code__.co_argcount == 2:
-            array = formula(population, period)
-        else:
-            array = formula(population, period, parameters_at)
+        try:
+            if formula.__code__.co_argcount == 2:
+                array = formula(population, period)
+            else:
+                array = formula(population, period, parameters_at)
+        finally:
+            if previous_tracing is not None:
+                parameters_at.set_tracing(*previous_tracing)
 
         return array
 
@@ -1459,6 +1504,17 @@ def get_branch(
         if self.trace:
             branch.trace = True
             branch.tracer = self.tracer
+            # The trace setter pointed the branch's parameter tree at a fresh
+            # tracer under the branch's name. Point it at the shared tracer
+            # again, and when the tree itself is shared, back at the caller,
+            # whose formula is still running and may read parameters before
+            # the branch calculates anything.
+            parameters = getattr(branch.tax_benefit_system, "parameters", None)
+            if parameters is not None:
+                if clone_system:
+                    parameters.set_tracing(self.tracer, name)
+                else:
+                    parameters.set_tracing(self.tracer, self.branch_name)
         return branch
 
     def derivative(