Share one population-aggregate list and one known-period walker (Route A follow-up)

docs/us-release-build-rule.md

@@ -161,8 +161,9 @@ A repeated pair is still refused. Source-level decisions fan out to the tail
 copy from the lowest surviving clone index. Row-level imputations treat the
 tail copy as a PUF-role row, as the base's own post-transfer stages do.
 Only a base without a raw spine ID may lack clone indices. An assembled table
-(one with a raw spine ID) must carry both its support channel and its clone
-index. `require_assembled_support_provenance` refuses one that lacks either.
+(one with a raw spine ID) must carry its support channel, its clone index and
+a complete, non-null `<entity>_source_id`. `require_assembled_support_provenance`
+refuses one that lacks any of them.
 The metadata-presence check includes the raw spine ID, so losing both columns
 cannot make an assembled table look historical. `support_role_series` calls
 the validator before resolving roles; this covers the prior-year and SSI
@@ -239,8 +240,9 @@ known violations at run time:
 - In a multi-batch pass, an engine refuses once it has computed a formula that
   aggregates over its whole simulation. In policyengine-us 2.2.1 these are the
   Medicaid SLCSP state sums behind `medicaid_cost` and the weighted household
-  and SPM-unit income deciles (`POST_EXPORT_POPULATION_AGGREGATE_VARIABLES`).
-  Each batch would compute them over itself alone.
+  and SPM-unit income deciles (`US_POPULATION_AGGREGATE_VARIABLES`, the same
+  list batched target materialization guards). Each batch would compute them
+  over itself alone.
 - A reform engine refuses once it has computed a formula that reads the
   baseline branch (`POST_EXPORT_BASELINE_BRANCH_READERS`). These are the
   Medicaid denominator a reform holds at baseline, and the behavioral-response

packages/microcosm-build/tests/test_us_batched_target_materialization.py

@@ -1640,3 +1640,27 @@ def test_population_aggregate_list_matches_installed_engine_sources() -> None:
     assert triaged_digest == (
         "f26eb560e474207fc1fa1d8828b62ba4791fb763089ea2ce7a85b57a259f2755"
     ), triaged_digest
+
+
+def test_refusal_reads_a_value_held_only_on_baseline() -> None:
+    """Target materialization and post-export scoring share one walker. It
+    also reads a ``reform=`` engine's ``baseline`` simulation, which holds its
+    own values and is not among the engine's branches."""
+    builder = _load_builder_module()
+
+    def engine(known, **extra):
+        return SimpleNamespace(
+            get_holder=lambda name: SimpleNamespace(
+                get_known_periods=lambda: known.get(name, [])
+            ),
+            branches={},
+            **extra,
+        )
+
+    root = engine({}, baseline=engine({"medicaid_slcsp_state_denominator": ["2024"]}))
+    with pytest.raises(
+        ValueError, match=r"computed medicaid_slcsp_state_denominator@2024"
+    ):
+        builder._refuse_batch_population_aggregates(
+            root, batch=1, batches=3, n_households=6
+        )

packages/microcosm-build/tests/test_us_post_export_scoring.py

@@ -745,14 +745,129 @@ def engine(known):
     branch = engine({"household_income_decile": ["2025"]})
     root.branches["mtr"] = branch
     branch.branches["loop"] = root  # a cycle is walked once
-    assert builder._post_export_known_periods(
-        root, builder.POST_EXPORT_POPULATION_AGGREGATE_VARIABLES
+    assert builder._engine_known_periods(
+        root, builder.US_POPULATION_AGGREGATE_VARIABLES
     ) == {
         ("medicaid_slcsp_state_denominator", "2024"),
         ("household_income_decile", "2025"),
     }
 
 
+_WALKER_VARIABLES = (
+    "household_income_decile",
+    "medicaid_slcsp_state_denominator",
+    "other",
+)
+
+
+def test_known_period_walker_matches_a_reference_traversal(builder) -> None:
+    """The shared walker returns exactly the known (variable, period) pairs of
+    every engine reachable from the root through live branches, recorded
+    detached branches and ``baseline``, for the watched variables only, and
+    terminates on cycles. Checked against an independent breadth-first walk
+    over random engine graphs. Hypothesis is a workspace dependency; the
+    wheels job installs no test extras, so it skips there."""
+    pytest.importorskip("hypothesis")
+    from hypothesis import given, settings
+    from hypothesis import strategies as st
+
+    @st.composite
+    def engine_graphs(draw):
+        n = draw(st.integers(1, 7))
+        known = [
+            {
+                variable: draw(st.lists(st.integers(2020, 2035), max_size=3))
+                for variable in _WALKER_VARIABLES
+            }
+            for _ in range(n)
+        ]
+        nodes = [
+            SimpleNamespace(
+                get_holder=(
+                    lambda name, k=k: SimpleNamespace(
+                        get_known_periods=lambda: k.get(name, [])
+                    )
+                ),
+                branches={},
+                _target_materialization_branches=[],
+                baseline=None,
+            )
+            for k in known
+        ]
+        edges: list[tuple[int, int]] = []
+        for i, node in enumerate(nodes):
+            for j in draw(st.lists(st.integers(0, n - 1), max_size=3)):
+                node.branches[f"b{j}"] = nodes[j]
+                edges.append((i, j))
+            for j in draw(st.lists(st.integers(0, n - 1), max_size=2)):
+                node._target_materialization_branches.append(nodes[j])
+                edges.append((i, j))
+            baseline = draw(st.one_of(st.none(), st.integers(0, n - 1)))
+            if baseline is not None:
+                node.baseline = nodes[baseline]
+                edges.append((i, baseline))
+        return nodes, known, edges
+
+    @settings(max_examples=200, deadline=None)
+    @given(graph=engine_graphs(), watched=st.sets(st.sampled_from(_WALKER_VARIABLES)))
+    def check(graph, watched):
+        nodes, known, edges = graph
+        reachable, frontier = {0}, [0]
+        while frontier:
+            current = frontier.pop()
+            for source, target in edges:
+                if source == current and target not in reachable:
+                    reachable.add(target)
+                    frontier.append(target)
+        expected = {
+            (variable, str(period))
+            for index in reachable
+            for variable in watched
+            for period in known[index][variable]
+        }
+        assert (
+            builder._engine_known_periods(nodes[0], tuple(sorted(watched))) == expected
+        )
+
+    check()
+
+
+def test_batch_invariance_check_reads_a_value_held_only_on_baseline(
+    builder,
+) -> None:
+    """Post-export scoring and target materialization share one walker. It
+    also reads a ``reform=`` engine's ``baseline`` simulation, which holds its
+    own values and is not among the engine's branches."""
+
+    def engine(known, **extra):
+        return SimpleNamespace(
+            get_holder=lambda name: SimpleNamespace(
+                get_known_periods=lambda: known.get(name, [])
+            ),
+            branches={},
+            **extra,
+        )
+
+    baseline = engine({"medicaid_slcsp_state_denominator": ["2024"]})
+    root = engine({}, baseline=baseline)
+    assert builder._engine_known_periods(
+        root, builder.US_POPULATION_AGGREGATE_VARIABLES
+    ) == {("medicaid_slcsp_state_denominator", "2024")}
+    with pytest.raises(
+        RuntimeError,
+        match=r"not batch-invariant.*medicaid_slcsp_state_denominator@2024 "
+        r"\(aggregates over",
+    ):
+        builder._assert_post_export_scoring_is_batch_invariant(
+            root,
+            builder.US_POPULATION_AGGREGATE_VARIABLES,
+            set(),
+            label="fixture",
+            batched=True,
+            reform=False,
+        )
+
+
 def test_a_reform_engine_refuses_a_baseline_branch_reader(builder, tmp_path) -> None:
     """A batch reform engine is built from the reform's system alone, so it
     has no baseline branch: a reform pass refuses a formula that reads one,
@@ -778,7 +893,7 @@ def test_a_reform_engine_refuses_a_baseline_branch_reader(builder, tmp_path) ->
     assert all(engine.dataset is None for engine in log.constructions)
     assert builder._post_export_watched_variables(batched=False, reform=False) == ()
     assert set(builder._post_export_watched_variables(batched=True, reform=True)) == {
-        *builder.POST_EXPORT_POPULATION_AGGREGATE_VARIABLES,
+        *builder.US_POPULATION_AGGREGATE_VARIABLES,
         *builder.POST_EXPORT_BASELINE_BRANCH_READERS,
     }
 
@@ -899,9 +1014,9 @@ def test_watched_formula_lists_match_the_installed_engine(builder) -> None:
     aggregates = _variables_reaching(
         modules, _POPULATION_AGGREGATE_MARKER, weight_reads=True
     )
-    assert sorted(aggregates) == sorted(
-        builder.POST_EXPORT_POPULATION_AGGREGATE_VARIABLES
-    ), aggregates
+    assert sorted(aggregates) == sorted(builder.US_POPULATION_AGGREGATE_VARIABLES), (
+        aggregates
+    )
     readers = _variables_reaching(modules, _BASELINE_BRANCH_MARKER)
     assert sorted(readers) == sorted(
         {

tools/build_us_fiscal_refresh_release.py

@@ -4317,8 +4317,8 @@ def _select_households_by_position(frame: Frame, positions: np.ndarray) -> Frame
 # Totals then differ from one whole-pool simulation only in floating-point
 # summation order, PROVIDED every scored measure is additive across
 # households: a formula that aggregates over its simulation's whole population
-# (``POST_EXPORT_POPULATION_AGGREGATE_VARIABLES``: the Medicaid SLCSP state
-# sums and the weighted income deciles) sees one batch, not the file. Every
+# (``US_POPULATION_AGGREGATE_VARIABLES``: the Medicaid SLCSP state sums and
+# the weighted income deciles) sees one batch, not the file. Every
 # batch engine of a multi-batch pass refuses once it has computed one of them,
 # so a measure that reaches such a formula fails the stage instead of scoring
 # a batch-local aggregate.
@@ -4331,16 +4331,22 @@ def _select_households_by_position(frame: Frame, positions: np.ndarray) -> Frame
 PostExportKey = tuple[str, int, str | None]
 
 #: policyengine-us formulas that aggregate over their simulation's whole
-#: population instead of within one household, so a batch engine computes them
-#: over its own batch only: ``medicaid_slcsp_state_average_cost_index`` and
+#: population instead of within one household, so a household-batch engine
+#: computes them over its own batch only. In policyengine-us 2.2.1,
+#: ``medicaid_slcsp_state_average_cost_index`` and
 #: ``medicaid_slcsp_state_denominator`` sum person weights by state
-#: (``sum_by_state``) and feed ``medicaid_cost`` (and through it ``medicaid``
-#: and ``household_health_benefits``), and the two income deciles are weighted
-#: ranks over every household or SPM unit. A multi-batch pass refuses an engine
-#: that computed any of them. ``test_us_post_export_scoring.py`` pins this list
-#: against the installed engine's variable sources and exercises the Medicaid
-#: refusal on a written fixture H5.
-POST_EXPORT_POPULATION_AGGREGATE_VARIABLES = (
+#: (``sum_by_state``) and feed ``medicaid_cost_if_enrolled``, and through it
+#: ``medicaid_cost``, ``medicaid`` and ``household_health_benefits``; the two
+#: income deciles are weighted ranks over every household or SPM unit. Batched
+#: target materialization (``_refuse_batch_population_aggregates``) and
+#: multi-batch post-export scoring
+#: (``_assert_post_export_scoring_is_batch_invariant``) refuse an engine that
+#: computed any of them. ``test_us_batched_target_materialization.py`` pins
+#: this list against the installed engine's variable sources; that pattern
+#: scan is a drift detector, not a proof of household locality
+#: (microcosm#956). ``test_us_post_export_scoring.py`` exercises the
+#: post-export Medicaid refusal on a written fixture H5.
+US_POPULATION_AGGREGATE_VARIABLES = (
     "household_income_decile",
     "medicaid_slcsp_state_average_cost_index",
     "medicaid_slcsp_state_denominator",
@@ -4353,7 +4359,7 @@ def _select_households_by_position(frame: Frame, positions: np.ndarray) -> Frame
 #: ``get_behavioral_response_measurements``, which measures a reform against
 #: ``get_branch("baseline")`` and would measure it against itself here. A reform
 #: pass refuses an engine that computed any of them, whatever its batch count.
-#: Pinned with the aggregates above.
+#: Pinned in ``test_us_post_export_scoring.py``.
 POST_EXPORT_BASELINE_BRANCH_READERS = (
     "income_elasticity_lsr",
     "medicaid_slcsp_state_denominator",
@@ -4548,41 +4554,12 @@ def calculate(self, key: PostExportKey):
 
 def _post_export_watched_variables(*, batched: bool, reform: bool) -> tuple[str, ...]:
     """The formulas a batch engine may not compute (see the block comment)."""
-    watched = POST_EXPORT_POPULATION_AGGREGATE_VARIABLES if batched else ()
+    watched = US_POPULATION_AGGREGATE_VARIABLES if batched else ()
     if reform:
         watched = (*watched, *POST_EXPORT_BASELINE_BRANCH_READERS)
     return tuple(dict.fromkeys(watched))
 
 
-def _post_export_known_periods(
-    simulation, variables: Sequence[str]
-) -> set[tuple[str, str]]:
-    """The (variable, period) values an engine or any live branch of it holds.
-
-    The holder's ``get_known_periods`` includes memory and disk storage.
-    Branches keep their own holders; the ones still attached are read too.
-    The real-engine guard test checks that the watched Medicaid computation
-    leaves a known period visible through this API.
-    """
-    known: set[tuple[str, str]] = set()
-    stack = [simulation]
-    seen: set[int] = set()
-    while stack:
-        current = stack.pop()
-        if id(current) in seen:
-            continue
-        seen.add(id(current))
-        for variable in variables:
-            known.update(
-                (variable, str(period))
-                for period in current.get_holder(variable).get_known_periods()
-            )
-        branches = getattr(current, "branches", None)
-        if isinstance(branches, dict):
-            stack.extend(branches.values())
-    return known
-
-
 def _assert_post_export_scoring_is_batch_invariant(
     simulation,
     watched: Sequence[str],
@@ -4597,13 +4574,13 @@ def _assert_post_export_scoring_is_batch_invariant(
     ``known_at_load`` holds what the engine held before scoring (values the
     written H5 stores as inputs, which batching does not change).
     """
-    computed = sorted(_post_export_known_periods(simulation, watched) - known_at_load)
+    computed = sorted(_engine_known_periods(simulation, watched) - known_at_load)
     if not computed:
         return
     reasons = []
     for variable, period in computed:
         why = []
-        if batched and variable in POST_EXPORT_POPULATION_AGGREGATE_VARIABLES:
+        if batched and variable in US_POPULATION_AGGREGATE_VARIABLES:
             why.append(
                 "aggregates over its simulation's whole population, here one "
                 "household batch"
@@ -5051,7 +5028,7 @@ def _score(
             with _automatic_gc_suspended():
                 simulation = self._construct(batch_frame, reform_system=reform_system)
                 try:
-                    known_at_load = _post_export_known_periods(simulation, watched)
+                    known_at_load = _engine_known_periods(simulation, watched)
                     engine = _AscendingPeriodEngine(simulation, label=label)
                     for key in keys:
                         parts[key].append(
@@ -6361,23 +6338,6 @@ def _base_simulation_household_columns(
     return columns
 
 
-#: policyengine-us formulas that aggregate over their simulation's whole
-#: population instead of within one household, so a household-batch engine
-#: computes them over its own batch only. In policyengine-us 2.2.1,
-#: ``medicaid_slcsp_state_average_cost_index`` and
-#: ``medicaid_slcsp_state_denominator`` sum person weights by state
-#: (``sum_by_state``) and feed ``medicaid_cost_if_enrolled``, and through it
-#: ``medicaid_cost`` and ``medicaid``; the two income deciles are weighted
-#: ranks over every household or SPM unit. The pattern scan pinning this list
-#: is a drift detector, not a proof of household locality (microcosm#956).
-US_POPULATION_AGGREGATE_VARIABLES = (
-    "household_income_decile",
-    "medicaid_slcsp_state_average_cost_index",
-    "medicaid_slcsp_state_denominator",
-    "spm_unit_income_decile",
-)
-
-
 @contextmanager
 def _record_engine_branches(simulation):
     """Retain branches created through get_branch until the batch is checked.
@@ -6438,7 +6398,17 @@ def retain(current):
 
 
 def _engine_known_periods(simulation, variables: Sequence[str]) -> set[tuple[str, str]]:
-    """Read holders in the engine, live branches and recorded detached branches."""
+    """The (variable, period) values an engine or any engine it reaches holds.
+
+    Batched target materialization and post-export scoring both read this. The
+    holder's ``get_known_periods`` includes memory and disk storage. The walk
+    covers live branches, branches target materialization recorded before they
+    were detached, and the ``baseline`` simulation a ``reform=`` engine keeps
+    (policyengine-core sets it only when a reform is passed, so the post-export
+    scorer's engines have none). The real-engine guard tests check that the
+    watched Medicaid computation leaves a known period visible through this
+    API.
+    """
 
     known: set[tuple[str, str]] = set()
     stack = [simulation]