From 66eb9caa8d4e5560eebef994f1c23bfb67268c4e Mon Sep 17 00:00:00 2001 From: PavelMakarchuk Date: Fri, 4 Sep 2026 15:11:16 +0200 Subject: [PATCH 1/6] Trace parameter reads in yearly formulas and scale reads; add dependency-map Two tracer gaps made FullTracer record almost no parameter reads for real models. The yearly ParameterNodeAtInstant was cached before the lazy tracing recast in _run_formula ran (defined_for and adds/subtracts evaluation build it first), so every parameters(period).gov.* read in a yearly formula went unrecorded while monthly formulas looked fine. Fixes #541. And TracingParameterNodeAtInstant only recorded scalar and array leaves, so scale and bracket reads (p.base.calc(age), p.max[children]) were invisible. Fixes #542. The trace setter now recasts the parameter tree (ParameterNode.set_tracing) and clears its at-instant caches whenever tracing switches on or off; _run_formula only follows the current branch. Scale children are recorded at their node with no scalar value. policyengine-core dependency-map -c builds on the fixed tracer: it traces the package's YAML tests (one per newly covered output variable per file) and optionally a microdata subsample, and writes readers (parameter path -> variables) and consumers (variable -> variables) with the package version, a fingerprint over its model surface, and the core version. Country packages need no wrapper. Co-Authored-By: Claude Fable 5.1 --- changelog.d/dependency-map.added.md | 1 + changelog.d/tracer-parameter-reads.fixed.md | 1 + .../parameters/parameter_node.py | 26 ++ .../scripts/policyengine_command.py | 36 ++ policyengine_core/simulations/simulation.py | 31 +- policyengine_core/tools/dependency_map.py | 309 ++++++++++++++++++ .../tracing_parameter_node_at_instant.py | 4 + tests/core/test_tracers.py | 50 +++ tests/core/tools/test_dependency_map.py | 109 ++++++ 9 files changed, 553 insertions(+), 14 deletions(-) create mode 100644 changelog.d/dependency-map.added.md create mode 100644 changelog.d/tracer-parameter-reads.fixed.md create mode 100644 policyengine_core/tools/dependency_map.py create mode 100644 tests/core/tools/test_dependency_map.py diff --git a/changelog.d/dependency-map.added.md b/changelog.d/dependency-map.added.md new file mode 100644 index 00000000..db81dc4c --- /dev/null +++ b/changelog.d/dependency-map.added.md @@ -0,0 +1 @@ +`policyengine-core dependency-map -c ` 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. diff --git a/changelog.d/tracer-parameter-reads.fixed.md b/changelog.d/tracer-parameter-reads.fixed.md new file mode 100644 index 00000000..a42ad588 --- /dev/null +++ b/changelog.d/tracer-parameter-reads.fixed.md @@ -0,0 +1 @@ +`FullTracer` now records parameter reads in yearly formulas: the parameter tree is recast for tracing when `simulation.trace` is set (and its at-instant caches cleared) instead of lazily inside `_run_formula`, after `defined_for` and `adds`/`subtracts` evaluation had already cached the yearly node untraced (#541). Scale and bracket reads (`p.rates.calc(x)`, `p.max[children]`) are recorded at their node too (#542). diff --git a/policyengine_core/parameters/parameter_node.py b/policyengine_core/parameters/parameter_node.py index 0df0e387..7d269960 100644 --- a/policyengine_core/parameters/parameter_node.py +++ b/policyengine_core/parameters/parameter_node.py @@ -234,6 +234,32 @@ def clear_parent_cache(self): if self.parent is not None: self.parent.clear_parent_cache() + def clear_at_instant_caches(self) -> None: + """Drop every cached at-instant node in this subtree. + + A cached ``ParameterNodeAtInstant`` is plain or tracing depending on + ``trace`` at the time it was built, so the caches must be cleared + whenever tracing is switched on or off. + """ + self._at_instant_cache.clear() + for child in self.children.values(): + clear = getattr(child, "clear_at_instant_caches", None) + if clear is not None: + clear() + else: + cache = getattr(child, "_at_instant_cache", None) + if cache is not None: + cache.clear() + + def set_tracing(self, tracer, branch_name: str) -> None: + """Route parameter reads through ``tracer`` (``None`` to stop).""" + was_tracing = self.trace + self.trace = tracer is not None + self.tracer = tracer + self.branch_name = branch_name + if self.trace != was_tracing: + self.clear_at_instant_caches() + def mark_as_modified(self): self.modified = True if self.parent is not None: diff --git a/policyengine_core/scripts/policyengine_command.py b/policyengine_core/scripts/policyengine_command.py index 7aa85c88..dbcab7c0 100644 --- a/policyengine_core/scripts/policyengine_command.py +++ b/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: /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()) diff --git a/policyengine_core/simulations/simulation.py b/policyengine_core/simulations/simulation.py index 0d4e001c..4fef05c6 100644 --- a/policyengine_core/simulations/simulation.py +++ b/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, whose + # at-instant nodes are cached as plain or tracing when first built. + # Recast (and clear) here rather than lazily in _run_formula: by the + # time a formula runs, defined_for and adds/subtracts evaluation has + # already cached the yearly node untraced. + 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(): @@ -1111,14 +1116,12 @@ def _run_formula( ) 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 + if self.trace and parameters_at is not None: + # Keep the parameter tree pointed at this simulation's tracer and + # branch (branch simulations share the tax-benefit system). The + # trace setter did the initial recast; this follows the branch. + 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 diff --git a/policyengine_core/tools/dependency_map.py b/policyengine_core/tools/dependency_map.py new file mode 100644 index 00000000..787e1e98 --- /dev/null +++ b/policyengine_core/tools/dependency_map.py @@ -0,0 +1,309 @@ +"""Trace the parameter → variable dependency map of a country package. + +Downstream tools need to know which variables a parameter path moves and +which variables feed which: validation matching against calibration targets +and scorecards, model-coverage audits, reform classifiers. Static scans of +formula source miss bracket, scale, and vectorised reads; the only exact +record is what the model reads at run time. This module runs simulations +under ``FullTracer`` and folds the trace trees into two edge sets: + + readers[parameter_path] -> variables whose formula read that parameter + consumers[variable] -> variables whose formula read that variable + +Parameter paths are recorded at the node the formula indexed, so a bracket +read ``p.base.calc(age)`` is recorded as ``gov.irs.credits.ctc.amount.base``. + +Populations +----------- +``tests`` builds a simulation for every YAML test under the package's + tests directory (skipping tests that apply reforms, + extensions, or inline parameter changes) and calculates the + test's outputs. Deterministic, no data download, and country + test suites deliberately exercise every program. +``microdata`` calculates every variable over a subsample of the package's + default microdata. Broad, but a formula behind ``defined_for`` + only runs when someone in the sample qualifies. +``both`` the union. + +Tests in one file mostly vary inputs for the same outputs and record the +same edges, so by default one test is kept per newly covered output +variable in each file; ``every_test`` traces them all. +""" + +from __future__ import annotations + +import hashlib +import importlib +import json +import sys +import time +from collections import defaultdict +from datetime import datetime, timezone +from importlib import metadata +from pathlib import Path +from typing import Callable, Iterable, Optional + +import yaml + +from policyengine_core.periods import ETERNITY +from policyengine_core.simulations import SimulationBuilder +from policyengine_core.tracers import FullTracer + +Edges = tuple[dict[str, set[str]], dict[str, set[str]]] +Progress = Callable[[str], None] + +DEFAULT_OUTPUT = Path("dependency-map.json") +# The model surface the map depends on; hashed into the fingerprint. +MODEL_SURFACE = ("entities.py", "parameters", "system.py", "variables") + + +def _quiet_tracer() -> None: + """Trace values are never read back; dropping them keeps memory flat.""" + if not getattr(FullTracer, "_dependency_map_quiet", False): + FullTracer.record_calculation_result = lambda self, value: None + FullTracer._dependency_map_quiet = True + + +def collect_edges(simulation, readers=None, consumers=None) -> Edges: + """Fold a traced simulation's trees into the two edge sets.""" + readers = defaultdict(set) if readers is None else readers + consumers = defaultdict(set) if consumers is None else consumers + seen: set[int] = set() + + def walk(node) -> None: + if id(node) in seen: + return + seen.add(id(node)) + for parameter in node.parameters: + readers[parameter.name].add(node.name) + for child in node.children: + if child.name != node.name: + consumers[child.name].add(node.name) + walk(child) + + for tree in simulation.tracer.trees: + walk(tree) + return readers, consumers + + +def iter_yaml_tests(paths: Iterable[Path], every_test: bool = False): + """Yield (file, test) for the tests worth tracing against the baseline.""" + for path in paths: + files = sorted(path.rglob("*.yaml")) if path.is_dir() else [path] + for file in files: + tests = yaml.safe_load(file.read_text()) or [] + if not isinstance(tests, list): + continue + covered: set[str] = set() + for test in tests: + if not isinstance(test, dict): + continue + inputs = test.get("input") or {} + if test.get("reforms") or test.get("extensions"): + continue + if any("." in key for key in inputs): + continue # inline parameter change: not the baseline system + outputs = test.get("output") or {} + if not outputs: + continue + if not every_test and covered >= set(outputs): + continue + covered |= set(outputs) + yield file, test + + +def trace_yaml_tests( + system, + paths: Iterable[Path], + progress: Optional[Progress] = None, + every_test: bool = False, +) -> tuple[Edges, dict[str, int]]: + _quiet_tracer() + readers: dict[str, set[str]] = defaultdict(set) + consumers: dict[str, set[str]] = defaultdict(set) + stats = {"tests": 0, "failed": 0} + for index, (file, test) in enumerate(iter_yaml_tests(paths, every_test)): + period = test.get("period") + try: + builder = SimulationBuilder() + builder.set_default_period(period) + simulation = builder.build_from_dict(system, test.get("input") or {}) + simulation.default_calculation_period = builder.default_period + simulation.trace = True + for output in test["output"]: + try: + simulation.calculate(output, period) + except Exception: # a failing test still traced what ran + pass + collect_edges(simulation, readers, consumers) + stats["tests"] += 1 + except Exception: # unbuildable situation: skip it + stats["failed"] += 1 + if progress and index % 500 == 0: + progress(f" {index} tests traced ({file.name})") + return (readers, consumers), stats + + +def trace_microdata( + microsimulation_class, + households: int = 2000, + year: int = 2026, + progress: Optional[Progress] = None, +) -> tuple[Edges, dict[str, int]]: + _quiet_tracer() + simulation = microsimulation_class() + simulation = simulation.subsample(n=households, seed=0) or simulation + simulation.trace = True + variables = simulation.tax_benefit_system.variables + stats = {"variables": 0, "failed": 0} + for index, name in enumerate(sorted(variables)): + variable = variables[name] + if variable.definition_period == ETERNITY: + periods = [ETERNITY] + elif variable.definition_period == "month": + periods = [f"{year}-01"] + else: + periods = [year, f"{year}-01"] + for period in periods: + try: + simulation.calculate(name, period) + stats["variables"] += 1 + break + except Exception: # any formula failure just skips the variable + continue + else: + stats["failed"] += 1 + if progress and index % 500 == 0: + progress(f" {index}/{len(variables)} variables") + return collect_edges(simulation), stats + + +def merge_edges(*edge_sets: Edges) -> Edges: + readers: dict[str, set[str]] = defaultdict(set) + consumers: dict[str, set[str]] = defaultdict(set) + for edge_readers, edge_consumers in edge_sets: + for path, names in edge_readers.items(): + readers[path] |= names + for name, users in edge_consumers.items(): + consumers[name] |= users + return readers, consumers + + +def model_fingerprint(package_root: Path) -> str: + """sha256 over the files the map depends on, so consumers can tell + whether a map still matches the model they run against.""" + digest = hashlib.sha256() + for relative in MODEL_SURFACE: + path = package_root / relative + files = ( + [path] + if path.is_file() + else sorted( + child + for child in path.rglob("*") + if child.is_file() + and "__pycache__" not in child.parts + and child.suffix not in {".pyc", ".pyo"} + ) + if path.is_dir() + else [] + ) + for file in files: + digest.update(file.relative_to(package_root).as_posix().encode()) + digest.update(b"\0") + digest.update(file.read_bytes()) + digest.update(b"\0") + return f"sha256:{digest.hexdigest()}" + + +def _package_version(package_name: str) -> Optional[str]: + for distribution in (package_name, package_name.replace("_", "-")): + try: + return metadata.version(distribution) + except metadata.PackageNotFoundError: + continue + return None + + +def build_dependency_map( + country_package_name: str, + population: str = "tests", + tests_root: Optional[Path] = None, + every_test: bool = False, + households: int = 2000, + year: int = 2026, + progress: Optional[Progress] = None, +) -> dict: + package = importlib.import_module(country_package_name) + package_root = Path(package.__file__).resolve().parent + tests_root = tests_root or package_root / "tests" + + edge_sets: list[Edges] = [] + populations: dict[str, dict] = {} + started = time.time() + if population in ("tests", "both"): + edges, stats = trace_yaml_tests( + package.CountryTaxBenefitSystem(), [tests_root], progress, every_test + ) + edge_sets.append(edges) + populations["tests"] = { + "root": str(tests_root), + "everyTest": every_test, + **stats, + } + if population in ("microdata", "both"): + edges, stats = trace_microdata( + package.Microsimulation, households, year, progress + ) + edge_sets.append(edges) + populations["microdata"] = {"households": households, "year": year, **stats} + if not edge_sets: + raise ValueError(f"unknown population {population!r}") + + readers, consumers = merge_edges(*edge_sets) + return { + "generatedAt": datetime.now(timezone.utc).isoformat(timespec="seconds"), + "model": { + "package": country_package_name, + "version": _package_version(country_package_name), + "fingerprint": model_fingerprint(package_root), + "coreVersion": metadata.version("policyengine-core"), + }, + "populations": populations, + "tracingSeconds": round(time.time() - started), + "readers": {path: sorted(names) for path, names in sorted(readers.items())}, + "consumers": {name: sorted(users) for name, users in sorted(consumers.items())}, + } + + +def write_dependency_map(payload: dict, output: Path = DEFAULT_OUTPUT) -> Path: + output.parent.mkdir(parents=True, exist_ok=True) + output.write_text(json.dumps(payload, separators=(",", ":")) + "\n") + return output + + +def main(parser) -> int: + args = parser.parse_args() + + def progress(message: str) -> None: + print(message, file=sys.stderr, flush=True) + + from policyengine_core.scripts import detect_country_package + + country_package_name = args.country_package or detect_country_package() + payload = build_dependency_map( + country_package_name, + population=args.population, + tests_root=Path(args.tests_root) if args.tests_root else None, + every_test=args.every_test, + households=args.households, + year=args.year, + progress=progress, + ) + output = write_dependency_map(payload, Path(args.output)) + progress( + f"wrote {output}: {len(payload['readers'])} parameter paths, " + f"{len(payload['consumers'])} consumed variables, " + f"{payload['tracingSeconds']}s" + ) + return 0 diff --git a/policyengine_core/tracers/tracing_parameter_node_at_instant.py b/policyengine_core/tracers/tracing_parameter_node_at_instant.py index d4bf1746..d1cd0a5c 100644 --- a/policyengine_core/tracers/tracing_parameter_node_at_instant.py +++ b/policyengine_core/tracers/tracing_parameter_node_at_instant.py @@ -76,5 +76,9 @@ def get_traced_child( if isinstance(child, (numpy.ndarray,) + parameters.ALLOWED_PARAM_TYPES): self.tracer.record_parameter_access(name, period, self.branch_name, child) + else: + # A scale or bracket (read through .calc() or [index]) is a + # parameter read too; record it at its node with no scalar value. + self.tracer.record_parameter_access(name, period, self.branch_name, None) return child diff --git a/tests/core/test_tracers.py b/tests/core/test_tracers.py index e653f39b..0fce8f7f 100644 --- a/tests/core/test_tracers.py +++ b/tests/core/test_tracers.py @@ -529,3 +529,53 @@ def test_browse_trace(): browsed_nodes = [node.name for node in tracer.browse_trace()] assert browsed_nodes == ["B", "C", "D", "E", "F"] + + +# Tests on parameter reads recorded through a simulation (#541, #542) +PARAMETER_PERIOD = "2017-01" + + +def _tree(simulation, name): + return next(tree for tree in simulation.tracer.trees if tree.name == name) + + +@mark.parametrize("simulation", [({"salary": 2000}, PARAMETER_PERIOD)], indirect=True) +def test_yearly_parameter_read_traced_after_node_was_cached(simulation): + # Reproduces #541: the at-instant node is cached untraced before tracing + # is switched on (defined_for / adds evaluation does this in practice). + simulation.tax_benefit_system.parameters("2017-01-01") + simulation.trace = True + + simulation.calculate("income_tax", PARAMETER_PERIOD) + + names = [p.name for p in _tree(simulation, "income_tax").parameters] + assert names == ["taxes.income_tax_rate"] + + +@mark.parametrize("simulation", [({"salary": 2000}, PARAMETER_PERIOD)], indirect=True) +def test_scale_read_is_traced(simulation): + # Reproduces #542: scale.calc(salary) is a parameter read. + simulation.trace = True + + simulation.calculate("social_security_contribution", PARAMETER_PERIOD) + + parameters = _tree(simulation, "social_security_contribution").parameters + assert [p.name for p in parameters] == ["taxes.social_security_contribution"] + assert parameters[0].value is None + + +@mark.parametrize("simulation", [({"salary": 2000}, PARAMETER_PERIOD)], indirect=True) +def test_switching_trace_off_restores_plain_parameter_nodes(simulation): + simulation.trace = True + assert isinstance( + simulation.tax_benefit_system.parameters("2017-01-01"), + TracingParameterNodeAtInstant, + ) + + simulation.trace = False + + assert not isinstance( + simulation.tax_benefit_system.parameters("2017-01-01"), + TracingParameterNodeAtInstant, + ) + simulation.calculate("income_tax", PARAMETER_PERIOD) diff --git a/tests/core/tools/test_dependency_map.py b/tests/core/tools/test_dependency_map.py new file mode 100644 index 00000000..f1f71bd9 --- /dev/null +++ b/tests/core/tools/test_dependency_map.py @@ -0,0 +1,109 @@ +import json +from pathlib import Path + +from policyengine_core.country_template import CountryTaxBenefitSystem +from policyengine_core.tools.dependency_map import ( + build_dependency_map, + iter_yaml_tests, + merge_edges, + model_fingerprint, + trace_yaml_tests, + write_dependency_map, +) + +import policyengine_core.country_template as country_template + +TESTS_ROOT = Path(country_template.__file__).resolve().parent / "tests" + + +def test_trace_yaml_tests_records_parameter_and_variable_edges(): + (readers, consumers), stats = trace_yaml_tests( + CountryTaxBenefitSystem(), [TESTS_ROOT / "income_tax.yaml"] + ) + + assert stats == {"tests": 1, "failed": 0} + assert readers["taxes.income_tax_rate"] == {"income_tax"} + assert consumers["salary"] == {"income_tax"} + + +def test_trace_yaml_tests_records_scale_reads(): + (readers, _), _ = trace_yaml_tests( + CountryTaxBenefitSystem(), [TESTS_ROOT / "social_security_contribution.yaml"] + ) + + assert readers["taxes.social_security_contribution"] == { + "social_security_contribution" + } + + +def test_iter_yaml_tests_keeps_one_test_per_new_output(tmp_path: Path): + (tmp_path / "cases.yaml").write_text( + """ +- name: plain + period: 2017-01 + input: {salary: 1000} + output: {income_tax: 150} +- name: with reform + period: 2017-01 + reforms: policyengine_core.country_template.reforms.some_reform + input: {salary: 1000} + output: {income_tax: 150} +- name: inline parameter change + period: 2017-01 + input: {salary: 1000, taxes.income_tax_rate: 0.5} + output: {income_tax: 500} +- name: no output + period: 2017-01 + input: {salary: 1000} +- name: same output again + period: 2017-01 + input: {salary: 2000} + output: {income_tax: 300} +- name: new output + period: 2017-01 + input: {salary: 2000} + output: {income_tax: 300, social_security_contribution: 40} +""" + ) + + names = [test["name"] for _, test in iter_yaml_tests([tmp_path])] + every = [test["name"] for _, test in iter_yaml_tests([tmp_path], every_test=True)] + + assert names == ["plain", "new output"] + assert every == ["plain", "same output again", "new output"] + + +def test_merge_edges_unions_both_sides(): + readers, consumers = merge_edges( + ({"gov.a": {"x"}}, {"x": {"y"}}), + ({"gov.a": {"z"}, "gov.b": {"w"}}, {"x": {"q"}}), + ) + + assert readers == {"gov.a": {"x", "z"}, "gov.b": {"w"}} + assert consumers == {"x": {"y", "q"}} + + +def test_model_fingerprint_changes_with_the_model_surface(tmp_path: Path): + (tmp_path / "parameters").mkdir() + (tmp_path / "parameters" / "rate.yaml").write_text("values: {2017-01-01: 0.1}") + before = model_fingerprint(tmp_path) + + (tmp_path / "parameters" / "rate.yaml").write_text("values: {2017-01-01: 0.2}") + + assert before.startswith("sha256:") + assert model_fingerprint(tmp_path) != before + + +def test_build_and_write_dependency_map_for_the_country_template(tmp_path: Path): + payload = build_dependency_map( + "policyengine_core.country_template", tests_root=TESTS_ROOT + ) + output = write_dependency_map(payload, tmp_path / "map.json") + + written = json.loads(output.read_text()) + assert written["model"]["package"] == "policyengine_core.country_template" + assert written["model"]["fingerprint"].startswith("sha256:") + assert written["model"]["coreVersion"] + assert written["populations"]["tests"]["tests"] > 0 + assert "income_tax" in written["readers"]["taxes.income_tax_rate"] + assert "income_tax" in written["consumers"]["salary"] From b3fc6304c56d3d60d333ac4f22e2de06d5a404e7 Mon Sep 17 00:00:00 2001 From: PavelMakarchuk Date: Fri, 4 Sep 2026 15:13:31 +0200 Subject: [PATCH 2/6] Skip gov.abolitions reads in the dependency map: core's neutralisation check reads them, not the formula Co-Authored-By: Claude Fable 5.1 --- policyengine_core/tools/dependency_map.py | 6 +++++- 1 file changed, 5 insertions(+), 1 deletion(-) diff --git a/policyengine_core/tools/dependency_map.py b/policyengine_core/tools/dependency_map.py index 787e1e98..eb98609c 100644 --- a/policyengine_core/tools/dependency_map.py +++ b/policyengine_core/tools/dependency_map.py @@ -53,6 +53,9 @@ Progress = Callable[[str], None] DEFAULT_OUTPUT = Path("dependency-map.json") +# Read by core's neutralisation check before a formula runs, not by the +# formula: a switch per variable, not a dependency. +IGNORED_PARAMETER_PREFIXES = ("gov.abolitions.",) # The model surface the map depends on; hashed into the fingerprint. MODEL_SURFACE = ("entities.py", "parameters", "system.py", "variables") @@ -75,7 +78,8 @@ def walk(node) -> None: return seen.add(id(node)) for parameter in node.parameters: - readers[parameter.name].add(node.name) + if not parameter.name.startswith(IGNORED_PARAMETER_PREFIXES): + readers[parameter.name].add(node.name) for child in node.children: if child.name != node.name: consumers[child.name].add(node.name) From cf547a4ac899353c18a5da8b12c6bfec55796828 Mon Sep 17 00:00:00 2001 From: PavelMakarchuk Date: Fri, 4 Sep 2026 15:18:42 +0200 Subject: [PATCH 3/6] Scope the dependency map's quiet tracer to its own simulations The tool patched FullTracer.record_calculation_result globally, which leaked into any later trace in the same process (seven tracer tests failed when run after the tool's). It now uses a FullTracer subclass on its own simulations, and ParameterNode.set_tracing invalidates the at-instant caches when the tracer instance changes, not only when the flag does: cached tracing nodes capture the tracer they were built with, which branch simulations (which swap in the parent's tracer) rely on too. Co-Authored-By: Claude Fable 5.1 --- .../parameters/parameter_node.py | 6 +++-- policyengine_core/tools/dependency_map.py | 22 ++++++++++++------- 2 files changed, 18 insertions(+), 10 deletions(-) diff --git a/policyengine_core/parameters/parameter_node.py b/policyengine_core/parameters/parameter_node.py index 7d269960..f94c360e 100644 --- a/policyengine_core/parameters/parameter_node.py +++ b/policyengine_core/parameters/parameter_node.py @@ -253,11 +253,13 @@ def clear_at_instant_caches(self) -> None: def set_tracing(self, tracer, branch_name: str) -> None: """Route parameter reads through ``tracer`` (``None`` to stop).""" - was_tracing = self.trace + # Cached at-instant nodes capture the tracer they were built with, + # so a different tracer instance also invalidates them. + changed = (tracer is not None) != self.trace or tracer is not self.tracer self.trace = tracer is not None self.tracer = tracer self.branch_name = branch_name - if self.trace != was_tracing: + if changed: self.clear_at_instant_caches() def mark_as_modified(self): diff --git a/policyengine_core/tools/dependency_map.py b/policyengine_core/tools/dependency_map.py index eb98609c..a09b6cec 100644 --- a/policyengine_core/tools/dependency_map.py +++ b/policyengine_core/tools/dependency_map.py @@ -60,11 +60,19 @@ MODEL_SURFACE = ("entities.py", "parameters", "system.py", "variables") -def _quiet_tracer() -> None: +class _QuietFullTracer(FullTracer): """Trace values are never read back; dropping them keeps memory flat.""" - if not getattr(FullTracer, "_dependency_map_quiet", False): - FullTracer.record_calculation_result = lambda self, value: None - FullTracer._dependency_map_quiet = True + + def record_calculation_result(self, value) -> None: + pass + + +def _enable_tracing(simulation) -> None: + simulation.trace = True + simulation.tracer = _QuietFullTracer() + simulation.tax_benefit_system.parameters.set_tracing( + simulation.tracer, simulation.branch_name + ) def collect_edges(simulation, readers=None, consumers=None) -> Edges: @@ -122,7 +130,6 @@ def trace_yaml_tests( progress: Optional[Progress] = None, every_test: bool = False, ) -> tuple[Edges, dict[str, int]]: - _quiet_tracer() readers: dict[str, set[str]] = defaultdict(set) consumers: dict[str, set[str]] = defaultdict(set) stats = {"tests": 0, "failed": 0} @@ -133,7 +140,7 @@ def trace_yaml_tests( builder.set_default_period(period) simulation = builder.build_from_dict(system, test.get("input") or {}) simulation.default_calculation_period = builder.default_period - simulation.trace = True + _enable_tracing(simulation) for output in test["output"]: try: simulation.calculate(output, period) @@ -154,10 +161,9 @@ def trace_microdata( year: int = 2026, progress: Optional[Progress] = None, ) -> tuple[Edges, dict[str, int]]: - _quiet_tracer() simulation = microsimulation_class() simulation = simulation.subsample(n=households, seed=0) or simulation - simulation.trace = True + _enable_tracing(simulation) variables = simulation.tax_benefit_system.variables stats = {"variables": 0, "failed": 0} for index, name in enumerate(sorted(variables)): From de2b4090533608c640900e31596199ac6de9198a Mon Sep 17 00:00:00 2001 From: PavelMakarchuk Date: Fri, 4 Sep 2026 17:21:21 +0200 Subject: [PATCH 4/6] Keep the parameter at-instant cache plain and wrap for tracing on exit Clearing the at-instant caches whenever tracing switched or the tracer changed rebuilt the whole parameter tree each time, about two seconds on policyengine-us. Branch simulations swap tracers per formula, and the API traces some household simulations and not others on one shared system, so that cost landed on every consumer of tracing. The cache now holds plain nodes only, and _get_at_instant wraps one in TracingParameterNodeAtInstant on the way out when tracing is on. The wrapper reads the tracer and branch name from the parameter root at access time, so nothing is ever invalidated: set_tracing is three attribute writes, and branch labelling of parameter reads is now the current branch rather than whichever built the cache. Co-Authored-By: Claude Fable 5.1 --- .../parameters/parameter_node.py | 48 +++++++------------ .../tracing_parameter_node_at_instant.py | 26 ++++++++-- tests/core/test_tracers.py | 26 ++++++++++ 3 files changed, 66 insertions(+), 34 deletions(-) diff --git a/policyengine_core/parameters/parameter_node.py b/policyengine_core/parameters/parameter_node.py index f94c360e..6f466e1b 100644 --- a/policyengine_core/parameters/parameter_node.py +++ b/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( - node_at_instant, self.tracer, self.branch_name + return TracingParameterNodeAtInstant( + node_at_instant, self.tracer, self.branch_name, tracing_root=self ) - 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,33 +236,17 @@ def clear_parent_cache(self): if self.parent is not None: self.parent.clear_parent_cache() - def clear_at_instant_caches(self) -> None: - """Drop every cached at-instant node in this subtree. + def set_tracing(self, tracer, branch_name: str) -> None: + """Route parameter reads through ``tracer`` (``None`` to stop). - A cached ``ParameterNodeAtInstant`` is plain or tracing depending on - ``trace`` at the time it was built, so the caches must be cleared - whenever tracing is switched on or off. + Cheap by design: the at-instant cache holds plain nodes and the + tracing wrapper reads ``tracer`` and ``branch_name`` from this node + at access time, so nothing is rebuilt when tracing is switched or + a branch simulation swaps tracers. """ - self._at_instant_cache.clear() - for child in self.children.values(): - clear = getattr(child, "clear_at_instant_caches", None) - if clear is not None: - clear() - else: - cache = getattr(child, "_at_instant_cache", None) - if cache is not None: - cache.clear() - - def set_tracing(self, tracer, branch_name: str) -> None: - """Route parameter reads through ``tracer`` (``None`` to stop).""" - # Cached at-instant nodes capture the tracer they were built with, - # so a different tracer instance also invalidates them. - changed = (tracer is not None) != self.trace or tracer is not self.tracer self.trace = tracer is not None self.tracer = tracer self.branch_name = branch_name - if changed: - self.clear_at_instant_caches() def mark_as_modified(self): self.modified = True diff --git a/policyengine_core/tracers/tracing_parameter_node_at_instant.py b/policyengine_core/tracers/tracing_parameter_node_at_instant.py index d1cd0a5c..999b679a 100644 --- a/policyengine_core/tracers/tracing_parameter_node_at_instant.py +++ b/policyengine_core/tracers/tracing_parameter_node_at_instant.py @@ -27,10 +27,28 @@ def __init__( parameter_node_at_instant: ParameterNode, tracer: tracers.FullTracer, branch_name: str, + tracing_root=None, ) -> None: self.parameter_node_at_instant = parameter_node_at_instant - self.tracer = tracer - self.branch_name = branch_name + self._tracer = tracer + self._branch_name = branch_name + # A cached wrapper outlives the simulation that built it: branch + # simulations share the parameter tree and swap tracers and branch + # names per formula. Reading them from the tree's root at access + # time keeps the cache valid across those swaps. + self._tracing_root = tracing_root + + @property + def tracer(self) -> tracers.FullTracer: + if self._tracing_root is not None: + return self._tracing_root.tracer + return self._tracer + + @property + def branch_name(self) -> str: + if self._tracing_root is not None: + return self._tracing_root.branch_name + return self._branch_name def __getattr__( self, @@ -60,7 +78,9 @@ def get_traced_child( parameters.VectorialParameterNodeAtInstant, ), ): - return TracingParameterNodeAtInstant(child, self.tracer, self.branch_name) + return TracingParameterNodeAtInstant( + child, self._tracer, self._branch_name, self._tracing_root + ) if not isinstance(key, str) or isinstance( self.parameter_node_at_instant, diff --git a/tests/core/test_tracers.py b/tests/core/test_tracers.py index 0fce8f7f..ab0aeffd 100644 --- a/tests/core/test_tracers.py +++ b/tests/core/test_tracers.py @@ -579,3 +579,29 @@ def test_switching_trace_off_restores_plain_parameter_nodes(simulation): TracingParameterNodeAtInstant, ) simulation.calculate("income_tax", PARAMETER_PERIOD) + + +@mark.parametrize("simulation", [({"salary": 2000}, PARAMETER_PERIOD)], indirect=True) +def test_swapping_tracer_keeps_cached_nodes_and_records_into_new_tracer( + simulation, +): + """Branch simulations share the parameter tree and swap in the parent's + tracer; that must neither rebuild the at-instant tree nor keep recording + into the old tracer.""" + simulation.trace = True + parameters = simulation.tax_benefit_system.parameters + cached = parameters("2017-01-01") + old_tracer = simulation.tracer + + new_tracer = FullTracer() + parameters.set_tracing(new_tracer, "branch") + + assert ( + parameters("2017-01-01").parameter_node_at_instant + is cached.parameter_node_at_instant + ) + new_tracer._enter_calculation("income_tax", PARAMETER_PERIOD) + cached.taxes.income_tax_rate + assert new_tracer.trees[0].parameters[0].name == "taxes.income_tax_rate" + assert new_tracer.trees[0].parameters[0].branch_name == "branch" + assert old_tracer.trees == [] From b916dc51e3b4c7a99a7c3f3b47f8c24363b750c3 Mon Sep 17 00:00:00 2001 From: PavelMakarchuk Date: Mon, 7 Sep 2026 21:49:43 +0200 Subject: [PATCH 5/6] Trace core-resolved parameter reads, entity-scoped test outputs, and keep branch labels with their formula Review follow-up on #543. - Parameter-backed adds/subtracts lists and uprating factors are resolved on the parameter tree directly, bypassing the formula's tracing wrapper. Record them on the variable's trace node (Simulation._record_parameter_read). - TracingParameterNodeAtInstant captures its tracer and branch name when created instead of reading the tree's mutable state at access time; _run_formula restores the tree's tracing state after the formula, and get_branch hands the tree back to the caller after the branch's trace setter re-pointed it. A nested branch calculation no longer relabels the parameter reads its caller makes afterwards. - The dependency map reads test outputs the way the YAML runner does: entity singular and plural keys and period-keyed values, so entity-scoped cases are traced instead of silently skipped. Output calculations that raise are counted (outputErrors) and printed, and the module documents the output schema and how the counts bound completeness. - Changelog fragment and trace-setter comment describe the wrap-on-access design rather than cache clearing. Co-Authored-By: Claude Fable 5.1 --- changelog.d/tracer-parameter-reads.fixed.md | 2 +- .../parameters/parameter_node.py | 10 +- policyengine_core/simulations/simulation.py | 81 +++++++-- policyengine_core/tools/dependency_map.py | 97 +++++++++-- .../tracing_parameter_node_at_instant.py | 30 +--- tests/core/test_tracers.py | 157 +++++++++++++++++- tests/core/tools/test_dependency_map.py | 70 +++++++- 7 files changed, 383 insertions(+), 64 deletions(-) diff --git a/changelog.d/tracer-parameter-reads.fixed.md b/changelog.d/tracer-parameter-reads.fixed.md index a42ad588..26ab3a35 100644 --- a/changelog.d/tracer-parameter-reads.fixed.md +++ b/changelog.d/tracer-parameter-reads.fixed.md @@ -1 +1 @@ -`FullTracer` now records parameter reads in yearly formulas: the parameter tree is recast for tracing when `simulation.trace` is set (and its at-instant caches cleared) instead of lazily inside `_run_formula`, after `defined_for` and `adds`/`subtracts` evaluation had already cached the yearly node untraced (#541). Scale and bracket reads (`p.rates.calc(x)`, `p.max[children]`) are recorded at their node too (#542). +`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. diff --git a/policyengine_core/parameters/parameter_node.py b/policyengine_core/parameters/parameter_node.py index 6f466e1b..aca36a40 100644 --- a/policyengine_core/parameters/parameter_node.py +++ b/policyengine_core/parameters/parameter_node.py @@ -224,7 +224,7 @@ def _get_at_instant(self, instant: Instant) -> ParameterNodeAtInstant: self._at_instant_cache[instant] = node_at_instant if self.trace: return TracingParameterNodeAtInstant( - node_at_instant, self.tracer, self.branch_name, tracing_root=self + node_at_instant, self.tracer, self.branch_name ) return node_at_instant @@ -239,10 +239,10 @@ def clear_parent_cache(self): 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 the - tracing wrapper reads ``tracer`` and ``branch_name`` from this node - at access time, so nothing is rebuilt when tracing is switched or - a branch simulation swaps tracers. + 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 diff --git a/policyengine_core/simulations/simulation.py b/policyengine_core/simulations/simulation.py index 4fef05c6..c4ef2504 100644 --- a/policyengine_core/simulations/simulation.py +++ b/policyengine_core/simulations/simulation.py @@ -537,11 +537,11 @@ def trace(self, trace: SimpleTracer) -> None: self.tracer = FullTracer() else: self.tracer = SimpleTracer() - # Parameter reads are recorded through the parameter tree, whose - # at-instant nodes are cached as plain or tracing when first built. - # Recast (and clear) here rather than lazily in _run_formula: by the - # time a formula runs, defined_for and adds/subtracts evaluation has - # already cached the yearly node untraced. + # 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 ) @@ -860,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: @@ -1039,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: @@ -1061,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 @@ -1075,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." @@ -1092,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: @@ -1109,7 +1140,11 @@ 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." @@ -1117,19 +1152,26 @@ def _run_formula( return values parameters_at = self.tax_benefit_system.parameters + previous_tracing = None if self.trace and parameters_at is not None: - # Keep the parameter tree pointed at this simulation's tracer and - # branch (branch simulations share the tax-benefit system). The - # trace setter did the initial recast; this follows the branch. + # 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 @@ -1462,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( diff --git a/policyengine_core/tools/dependency_map.py b/policyengine_core/tools/dependency_map.py index a09b6cec..1fd475e3 100644 --- a/policyengine_core/tools/dependency_map.py +++ b/policyengine_core/tools/dependency_map.py @@ -27,7 +27,36 @@ Tests in one file mostly vary inputs for the same outputs and record the same edges, so by default one test is kept per newly covered output -variable in each file; ``every_test`` traces them all. +variable in each file; ``every_test`` traces them all. A test's ``output`` +block is read the way the YAML runner reads it: 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. + +Output +------ +The map is one JSON object:: + + generatedAt ISO-8601 UTC timestamp + model package, version, fingerprint (sha256 over the model + surface: entities.py, parameters/, system.py, + variables/), coreVersion + populations per population traced, its selection and counts: + tests: root, everyTest, tests (situations built), + failed (situations that could not be built), outputs + (variables calculated), outputErrors (output + calculations that raised; their edges are whatever + ran before the error); microdata: households, year, + variables (calculated), failed (raised at every period) + tracingSeconds wall time + readers parameter path -> sorted variables that read it + consumers variable -> sorted variables whose formula read it + +The map records what ran. A test that could not be built contributes no +edges, and an output that raised contributes only the edges recorded before +it raised, so ``failed`` and ``outputErrors`` bound how incomplete a map may +be; the command prints both, and a map with non-zero counts should be read +as a lower bound on the model's dependencies. """ from __future__ import annotations @@ -98,7 +127,42 @@ def walk(node) -> None: return readers, consumers -def iter_yaml_tests(paths: Iterable[Path], every_test: bool = False): +def iter_output_variables(output, system=None): + """Yield ``(variable, period)`` for each variable a test's output checks. + + Mirrors ``YamlItem.check_output``: a key is a variable, an entity's + singular key holding variables, or an entity's plural key holding + instances of variables. A dict value keyed by period checks the variable + at each of those periods; ``period`` is ``None`` for the test's own + period. Without a ``system`` every key is taken to be a variable. + """ + if not isinstance(output, dict): + return + singular = {} if system is None else {e.key for e in system.entities} + plural = {} if system is None else {e.plural for e in system.entities} + + def variable(name, value): + if isinstance(value, dict): + for period in value: + yield name, period + else: + yield name, None + + for key, value in output.items(): + if system is None or system.get_variable(key) is not None: + yield from variable(key, value) + elif key in singular and isinstance(value, dict): + for name, nested in value.items(): + yield from variable(name, nested) + elif key in plural and isinstance(value, dict): + for instance in value.values(): + if isinstance(instance, dict): + for name, nested in instance.items(): + yield from variable(name, nested) + # Any other key fails the YAML runner as an unknown variable. + + +def iter_yaml_tests(paths: Iterable[Path], every_test: bool = False, system=None): """Yield (file, test) for the tests worth tracing against the baseline.""" for path in paths: files = sorted(path.rglob("*.yaml")) if path.is_dir() else [path] @@ -115,12 +179,15 @@ def iter_yaml_tests(paths: Iterable[Path], every_test: bool = False): continue if any("." in key for key in inputs): continue # inline parameter change: not the baseline system - outputs = test.get("output") or {} + outputs = { + name + for name, _ in iter_output_variables(test.get("output"), system) + } if not outputs: continue - if not every_test and covered >= set(outputs): + if not every_test and covered >= outputs: continue - covered |= set(outputs) + covered |= outputs yield file, test @@ -132,8 +199,8 @@ def trace_yaml_tests( ) -> tuple[Edges, dict[str, int]]: readers: dict[str, set[str]] = defaultdict(set) consumers: dict[str, set[str]] = defaultdict(set) - stats = {"tests": 0, "failed": 0} - for index, (file, test) in enumerate(iter_yaml_tests(paths, every_test)): + stats = {"tests": 0, "failed": 0, "outputs": 0, "outputErrors": 0} + for index, (file, test) in enumerate(iter_yaml_tests(paths, every_test, system)): period = test.get("period") try: builder = SimulationBuilder() @@ -141,11 +208,12 @@ def trace_yaml_tests( simulation = builder.build_from_dict(system, test.get("input") or {}) simulation.default_calculation_period = builder.default_period _enable_tracing(simulation) - for output in test["output"]: + for name, output_period in iter_output_variables(test["output"], system): try: - simulation.calculate(output, period) - except Exception: # a failing test still traced what ran - pass + simulation.calculate(name, output_period or period) + stats["outputs"] += 1 + except Exception: # what ran before the error is still traced + stats["outputErrors"] += 1 collect_edges(simulation, readers, consumers) stats["tests"] += 1 except Exception: # unbuildable situation: skip it @@ -311,6 +379,13 @@ def progress(message: str) -> None: progress=progress, ) output = write_dependency_map(payload, Path(args.output)) + for name, counts in payload["populations"].items(): + summary = ", ".join( + f"{key}: {value}" + for key, value in counts.items() + if key not in ("root", "everyTest") + ) + progress(f"{name}: {summary}") progress( f"wrote {output}: {len(payload['readers'])} parameter paths, " f"{len(payload['consumers'])} consumed variables, " diff --git a/policyengine_core/tracers/tracing_parameter_node_at_instant.py b/policyengine_core/tracers/tracing_parameter_node_at_instant.py index 999b679a..f03e3ba2 100644 --- a/policyengine_core/tracers/tracing_parameter_node_at_instant.py +++ b/policyengine_core/tracers/tracing_parameter_node_at_instant.py @@ -27,28 +27,14 @@ def __init__( parameter_node_at_instant: ParameterNode, tracer: tracers.FullTracer, branch_name: str, - tracing_root=None, ) -> None: self.parameter_node_at_instant = parameter_node_at_instant - self._tracer = tracer - self._branch_name = branch_name - # A cached wrapper outlives the simulation that built it: branch - # simulations share the parameter tree and swap tracers and branch - # names per formula. Reading them from the tree's root at access - # time keeps the cache valid across those swaps. - self._tracing_root = tracing_root - - @property - def tracer(self) -> tracers.FullTracer: - if self._tracing_root is not None: - return self._tracing_root.tracer - return self._tracer - - @property - def branch_name(self) -> str: - if self._tracing_root is not None: - return self._tracing_root.branch_name - return self._branch_name + # Captured once: a wrapper belongs to the formula that obtained it. + # The parameter tree's own tracer and branch name move while a + # nested branch calculation runs; reading them at access time would + # label this formula's later reads with that branch (#543 review). + self.tracer = tracer + self.branch_name = branch_name def __getattr__( self, @@ -78,9 +64,7 @@ def get_traced_child( parameters.VectorialParameterNodeAtInstant, ), ): - return TracingParameterNodeAtInstant( - child, self._tracer, self._branch_name, self._tracing_root - ) + return TracingParameterNodeAtInstant(child, self.tracer, self.branch_name) if not isinstance(key, str) or isinstance( self.parameter_node_at_instant, diff --git a/tests/core/test_tracers.py b/tests/core/test_tracers.py index ab0aeffd..aa99129f 100644 --- a/tests/core/test_tracers.py +++ b/tests/core/test_tracers.py @@ -7,9 +7,12 @@ import numpy as np from pytest import approx, fixture, mark, raises +from policyengine_core.country_template.entities import Person from policyengine_core.country_template.variables.housing import ( HousingOccupancyStatus, ) +from policyengine_core.parameters import Parameter +from policyengine_core.periods import MONTH from policyengine_core.simulations import CycleError, Simulation, SpiralError from policyengine_core.tracers import ( FullTracer, @@ -17,6 +20,7 @@ TraceNode, TracingParameterNodeAtInstant, ) +from policyengine_core.variables import Variable from .parameters_fancy_indexing.test_fancy_indexing import parameters @@ -586,22 +590,157 @@ def test_swapping_tracer_keeps_cached_nodes_and_records_into_new_tracer( simulation, ): """Branch simulations share the parameter tree and swap in the parent's - tracer; that must neither rebuild the at-instant tree nor keep recording - into the old tracer.""" + tracer; that must not rebuild the at-instant tree. A wrapper keeps the + tracer and branch it was obtained under; the next access after the swap + records into the new tracer.""" simulation.trace = True parameters = simulation.tax_benefit_system.parameters - cached = parameters("2017-01-01") + before = parameters("2017-01-01") old_tracer = simulation.tracer new_tracer = FullTracer() parameters.set_tracing(new_tracer, "branch") + after = parameters("2017-01-01") - assert ( - parameters("2017-01-01").parameter_node_at_instant - is cached.parameter_node_at_instant - ) + assert after.parameter_node_at_instant is before.parameter_node_at_instant new_tracer._enter_calculation("income_tax", PARAMETER_PERIOD) - cached.taxes.income_tax_rate + after.taxes.income_tax_rate assert new_tracer.trees[0].parameters[0].name == "taxes.income_tax_rate" assert new_tracer.trees[0].parameters[0].branch_name == "branch" - assert old_tracer.trees == [] + old_tracer._enter_calculation("income_tax", PARAMETER_PERIOD) + before.taxes.income_tax_rate + assert old_tracer.trees[0].parameters[0].branch_name == "default" + assert len(new_tracer.trees[0].parameters) == 1 + + +# Reads the core makes on a variable's behalf, and branch labelling + + +class salary_plus_basic_income(Variable): + value_type = float + entity = Person + definition_period = MONTH + label = "adds: a variable and a parameter path" + adds = ["salary", "benefits.basic_income"] + + +class income_components_total(Variable): + value_type = float + entity = Person + definition_period = MONTH + label = "adds: a parameter holding the list" + adds = "taxes.income_components" + + +class minus_basic_income(Variable): + value_type = float + entity = Person + definition_period = MONTH + label = "subtracts: a parameter path" + subtracts = ["benefits.basic_income"] + + +class uprated_amount(Variable): + value_type = float + entity = Person + definition_period = MONTH + label = "no formula; carried forward by an uprating parameter" + uprating = "taxes.income_tax_rate" + + +class reads_around_a_branch(Variable): + value_type = float + entity = Person + definition_period = MONTH + label = "reads a parameter before and after a nested branch calculation" + + def formula(person, period, parameters): + before = parameters(period).taxes.income_tax_rate + person.simulation.get_branch("inner").calculate("income_tax", period) + after = parameters(period).taxes.income_tax_rate + return before + after + person("salary", period) * 0 + + +def _parameter_reads(simulation, name): + return [ + (p.name, p.period, p.branch_name, p.value) + for p in _tree(simulation, name).parameters + ] + + +def test_parameter_backed_adds_and_subtracts_are_traced( + tax_benefit_system, make_simulation +): + tax_benefit_system.parameters.taxes.add_child( + "income_components", + Parameter( + "taxes.income_components", + {"values": {"2015-01-01": {"value": ["salary"]}}}, + ), + ) + for variable in ( + salary_plus_basic_income, + income_components_total, + minus_basic_income, + ): + tax_benefit_system.add_variable(variable) + simulation = make_simulation(tax_benefit_system, {"salary": 2000}, PARAMETER_PERIOD) + simulation.trace = True + basic_income = tax_benefit_system.parameters.benefits.basic_income("2017-01-01") + + assert simulation.calculate("salary_plus_basic_income", PARAMETER_PERIOD) == approx( + 2000 + basic_income + ) + assert simulation.calculate("income_components_total", PARAMETER_PERIOD) == 2000 + assert simulation.calculate("minus_basic_income", PARAMETER_PERIOD) == approx( + -basic_income + ) + + assert _parameter_reads(simulation, "salary_plus_basic_income") == [ + ("benefits.basic_income", "2017-01-01", "default", basic_income) + ] + assert _parameter_reads(simulation, "income_components_total") == [ + ("taxes.income_components", "2017-01-01", "default", ["salary"]) + ] + assert _parameter_reads(simulation, "minus_basic_income") == [ + ("benefits.basic_income", "2017-01-01", "default", basic_income) + ] + + +def test_uprating_parameter_reads_are_traced(tax_benefit_system, make_simulation): + tax_benefit_system.add_variable(uprated_amount) + simulation = make_simulation(tax_benefit_system, {"salary": 2000}, "2017-01") + simulation.set_input("uprated_amount", "2017-01", [100]) + simulation.trace = True + rate = tax_benefit_system.parameters.taxes.income_tax_rate + + assert simulation.calculate("uprated_amount", "2017-02") == approx( + 100 * rate("2017-02-01") / rate("2017-01-01") + ) + + assert _parameter_reads(simulation, "uprated_amount") == [ + ("taxes.income_tax_rate", "2017-01-01", "default", rate("2017-01-01")), + ("taxes.income_tax_rate", "2017-02-01", "default", rate("2017-02-01")), + ] + + +def test_nested_branch_does_not_relabel_the_callers_parameter_reads( + tax_benefit_system, make_simulation +): + tax_benefit_system.add_variable(reads_around_a_branch) + simulation = make_simulation(tax_benefit_system, {"salary": 2000}, PARAMETER_PERIOD) + simulation.trace = True + + simulation.calculate("reads_around_a_branch", PARAMETER_PERIOD) + + outer = _tree(simulation, "reads_around_a_branch") + assert [(p.name, p.branch_name) for p in outer.parameters] == [ + ("taxes.income_tax_rate", "default"), + ("taxes.income_tax_rate", "default"), + ] + inner = next(child for child in outer.children if child.name == "income_tax") + assert [(p.name, p.branch_name) for p in inner.parameters] == [ + ("taxes.income_tax_rate", "inner") + ] + parameters = tax_benefit_system.parameters + assert (parameters.tracer, parameters.branch_name) == (simulation.tracer, "default") diff --git a/tests/core/tools/test_dependency_map.py b/tests/core/tools/test_dependency_map.py index f1f71bd9..50964f6b 100644 --- a/tests/core/tools/test_dependency_map.py +++ b/tests/core/tools/test_dependency_map.py @@ -21,11 +21,79 @@ def test_trace_yaml_tests_records_parameter_and_variable_edges(): CountryTaxBenefitSystem(), [TESTS_ROOT / "income_tax.yaml"] ) - assert stats == {"tests": 1, "failed": 0} + assert stats == {"tests": 1, "failed": 0, "outputs": 1, "outputErrors": 0} assert readers["taxes.income_tax_rate"] == {"income_tax"} assert consumers["salary"] == {"income_tax"} +def test_trace_yaml_tests_follows_entity_scoped_and_period_keyed_outputs( + tmp_path: Path, +): + (tmp_path / "cases.yaml").write_text( + """ +- name: plural entity outputs + period: 2017-01 + input: + persons: + alice: {salary: 1000} + households: + home: {parents: [alice]} + output: + persons: + alice: {income_tax: 150} +- name: singular entity output keyed by period + period: 2017-01 + input: {salary: 1000} + output: + person: + social_security_contribution: {2017-01: 20} +- name: same variable again under the plural key + period: 2017-01 + input: + persons: + bob: {salary: 3000} + households: + home: {parents: [bob]} + output: + persons: + bob: {income_tax: 450} +""" + ) + system = CountryTaxBenefitSystem() + + names = [test["name"] for _, test in iter_yaml_tests([tmp_path], system=system)] + (readers, consumers), stats = trace_yaml_tests(system, [tmp_path]) + + assert names == ["plural entity outputs", "singular entity output keyed by period"] + assert stats == {"tests": 2, "failed": 0, "outputs": 2, "outputErrors": 0} + assert readers["taxes.income_tax_rate"] == {"income_tax"} + assert readers["taxes.social_security_contribution"] == { + "social_security_contribution" + } + assert consumers["salary"] == {"income_tax", "social_security_contribution"} + + +def test_trace_yaml_tests_counts_outputs_that_fail_to_calculate(tmp_path: Path): + (tmp_path / "cases.yaml").write_text( + """ +- name: an output the model does not define + period: 2017-01 + input: {salary: 1000} + output: + person: + no_such_variable: 1 + social_security_contribution: 20 +""" + ) + + (readers, _), stats = trace_yaml_tests(CountryTaxBenefitSystem(), [tmp_path]) + + assert stats == {"tests": 1, "failed": 0, "outputs": 1, "outputErrors": 1} + assert readers["taxes.social_security_contribution"] == { + "social_security_contribution" + } + + def test_trace_yaml_tests_records_scale_reads(): (readers, _), _ = trace_yaml_tests( CountryTaxBenefitSystem(), [TESTS_ROOT / "social_security_contribution.yaml"] From f140bc27a6c018d359d1038c0fffcff6ddd74a37 Mon Sep 17 00:00:00 2001 From: PavelMakarchuk Date: Wed, 9 Sep 2026 17:19:40 +0200 Subject: [PATCH 6/6] Trace structural reads of a parameter node and describe scales in traces Review of #543, second round. Formulas read a node's structure in four places in policyengine-us: iterating p.node._children, taking its keys, taking its values and reading a nested node, and len(). The wrapper treated the _children dict like a scale, recording a path ending in ._children with no value, and handed out unwrapped nodes whose later reads went unrecorded. TracingChildren is a mapping view that records a read of the node's structure at the node's canonical path, with the child names as the value, and returns children through the parent's tracing so nested nodes stay wrapped. The wrapper also iterates and answers membership the same way. On the four formulas the map now records gov.usda.snap.work_requirements.abawd.waived_counties and its state children, gov.states.ut.dwf.ccap.copay.income_limits and the nested household-size node, gov.states.co.ccap.entry.fpg_rate, and gov.aca.la_county_rating_area, and nothing ending in ._children. Tax scales were recorded with a None value, which is enough for the dependency map but leaves serialized traces without the schedule. A scale read now records its class, thresholds, and rates or amounts as plain lists, so get_serialized_flat_trace carries the brackets. Co-Authored-By: Claude Fable 5.1 --- .../tracing_parameter_node_at_instant.py | 85 ++++++++++++++++++- tests/core/test_tracers.py | 54 +++++++++++- 2 files changed, 136 insertions(+), 3 deletions(-) diff --git a/policyengine_core/tracers/tracing_parameter_node_at_instant.py b/policyengine_core/tracers/tracing_parameter_node_at_instant.py index f03e3ba2..2c117123 100644 --- a/policyengine_core/tracers/tracing_parameter_node_at_instant.py +++ b/policyengine_core/tracers/tracing_parameter_node_at_instant.py @@ -1,11 +1,13 @@ from __future__ import annotations import typing +from collections.abc import Iterator, Mapping from typing import Union import numpy from policyengine_core import parameters +from policyengine_core.taxscales import TaxScaleLike from .. import tracers @@ -40,9 +42,35 @@ def __getattr__( self, key: str, ) -> Union[TracingParameterNodeAtInstant, Child]: + if key == "_children": + # A formula reading the node's structure (iterating child names, + # picking a child by computed key). Keep it traced rather than + # handing out the raw dict of unwrapped nodes. + return TracingChildren(self) child = getattr(self.parameter_node_at_instant, key) return self.get_traced_child(child, key) + def _child_names(self) -> list[str]: + return list(self.parameter_node_at_instant._children) + + def _record_structure_read(self) -> None: + """A read of which children a node has: recorded at the node's own + path, with the child names as its value.""" + self.tracer.record_parameter_access( + self.parameter_node_at_instant._name, + self.parameter_node_at_instant._instant_str, + self.branch_name, + self._child_names(), + ) + + def __iter__(self) -> Iterator[str]: + self._record_structure_read() + return iter(self._child_names()) + + def __contains__(self, key: object) -> bool: + self._record_structure_read() + return key in self.parameter_node_at_instant._children + def __getitem__( self, key: str, @@ -80,9 +108,62 @@ def get_traced_child( if isinstance(child, (numpy.ndarray,) + parameters.ALLOWED_PARAM_TYPES): self.tracer.record_parameter_access(name, period, self.branch_name, child) + elif isinstance(child, TaxScaleLike): + # A scale (read through .calc() or [index]) is a parameter read + # too; record its brackets, JSON-safe, so a serialized trace + # keeps the schedule and not just its path. + self.tracer.record_parameter_access( + name, period, self.branch_name, describe_tax_scale(child) + ) else: - # A scale or bracket (read through .calc() or [index]) is a - # parameter read too; record it at its node with no scalar value. self.tracer.record_parameter_access(name, period, self.branch_name, None) return child + + +def describe_tax_scale(scale: TaxScaleLike) -> dict: + """The schedule as plain lists: its class, thresholds, and the rates + or amounts it applies. Values are converted so the flat trace can be + serialized without a NumPy-aware encoder.""" + description: dict = {"type": type(scale).__name__} + for attribute in ("thresholds", "rates", "amounts"): + values = getattr(scale, attribute, None) + if values is not None: + description[attribute] = [_plain(v) for v in values] + return description + + +def _plain(value): + if isinstance(value, numpy.generic): + return value.item() + return value + + +class TracingChildren(Mapping): + """The child mapping of a traced node. + + Iterating, counting, or testing membership records a read of the + node's structure at the node's canonical path (never a path ending in + ``._children``, which nothing can resolve). Looking a child up returns + it through the parent's tracing, so nested nodes stay wrapped and + their later reads are recorded. + """ + + def __init__(self, parent: TracingParameterNodeAtInstant) -> None: + self._parent = parent + + def __getitem__(self, key: str): + child = self._parent.parameter_node_at_instant._children[key] + return self._parent.get_traced_child(child, key) + + def __iter__(self) -> Iterator[str]: + self._parent._record_structure_read() + return iter(self._parent._child_names()) + + def __len__(self) -> int: + self._parent._record_structure_read() + return len(self._parent._child_names()) + + def __contains__(self, key: object) -> bool: + self._parent._record_structure_read() + return key in self._parent.parameter_node_at_instant._children diff --git a/tests/core/test_tracers.py b/tests/core/test_tracers.py index aa99129f..11566226 100644 --- a/tests/core/test_tracers.py +++ b/tests/core/test_tracers.py @@ -565,7 +565,7 @@ def test_scale_read_is_traced(simulation): parameters = _tree(simulation, "social_security_contribution").parameters assert [p.name for p in parameters] == ["taxes.social_security_contribution"] - assert parameters[0].value is None + assert parameters[0].value["type"] == "MarginalRateTaxScale" @mark.parametrize("simulation", [({"salary": 2000}, PARAMETER_PERIOD)], indirect=True) @@ -744,3 +744,55 @@ def test_nested_branch_does_not_relabel_the_callers_parameter_reads( ] parameters = tax_benefit_system.parameters assert (parameters.tracer, parameters.branch_name) == (simulation.tracer, "default") + + +@mark.parametrize("simulation", [({"salary": 2000}, PARAMETER_PERIOD)], indirect=True) +def test_reading_a_nodes_children_is_recorded_at_the_node_and_keeps_them_wrapped( + simulation, +): + """Formulas iterate a node's children or pick one by computed key. That + must record the node's own path (never one ending in ._children) with + the names read, and hand back children that are still traced.""" + simulation.trace = True + tracer = simulation.tracer + tracer._enter_calculation("waived_area", PARAMETER_PERIOD) + taxes = simulation.tax_benefit_system.parameters("2017-01-01").taxes + + names = list(taxes._children) + assert len(taxes._children) == len(names) + assert "housing_tax" in taxes._children + housing = next(iter(taxes._children.values())) + housing.rate + + recorded = [(p.name, p.value) for p in tracer.trees[0].parameters] + assert ("taxes", names) in recorded + assert not any(name.endswith("._children") for name, _ in recorded) + assert isinstance(housing, TracingParameterNodeAtInstant) + assert "taxes.housing_tax.rate" in [name for name, _ in recorded] + + +@mark.parametrize("simulation", [({"salary": 2000}, PARAMETER_PERIOD)], indirect=True) +def test_iterating_a_traced_node_yields_child_names(simulation): + simulation.trace = True + simulation.tracer._enter_calculation("x", PARAMETER_PERIOD) + taxes = simulation.tax_benefit_system.parameters("2017-01-01").taxes + + assert sorted(taxes) == sorted(taxes.parameter_node_at_instant._children) + + +@mark.parametrize("simulation", [({"salary": 2000}, PARAMETER_PERIOD)], indirect=True) +def test_a_scale_read_records_its_brackets(simulation): + simulation.trace = True + tracer = simulation.tracer + tracer._enter_calculation("social_security_contribution", PARAMETER_PERIOD) + + simulation.tax_benefit_system.parameters( + "2017-01-01" + ).taxes.social_security_contribution + + recorded = tracer.trees[0].parameters[0] + assert recorded.name == "taxes.social_security_contribution" + assert recorded.value["type"] == "MarginalRateTaxScale" + assert recorded.value["thresholds"] == [0.0, 6000.0, 12400.0] + assert recorded.value["rates"] == [0.02, 0.06, 0.12] + json.dumps(tracer.get_serialized_flat_trace())