Declare local H5 postprocessor dependencies

Author: anth-volk

policyengine_us_data/build_outputs/builder.py

@@ -22,6 +22,7 @@
     "LocalAreaDatasetBuilder",
     "PayloadPostProcessor",
     "PayloadPostProcessorRun",
+    "PayloadPostProcessorSpec",
     "PayloadPostProcessorResult",
 ]
 
@@ -38,6 +39,8 @@ class PayloadPostProcessorResult(Protocol):
 class PayloadPostProcessor(Protocol):
     """Country- or product-specific processor for an H5 payload."""
 
+    spec: PayloadPostProcessorSpec
+
     def apply(
         self,
         *,
@@ -47,10 +50,19 @@ def apply(
         """Return a processed `H5Payload` or structured result with `.payload`."""
 
 
+@dataclass(frozen=True)
+class PayloadPostProcessorSpec:
+    """Stable identity and ordering requirements for a payload postprocessor."""
+
+    key: str
+    requires: tuple[str, ...] = ()
+
+
 @dataclass(frozen=True)
 class PayloadPostProcessorRun:
     """Result metadata for one payload postprocessor invocation."""
 
+    key: str
     name: str
     postprocessor_type: type
     result: PostProcessorReturn
@@ -97,7 +109,7 @@ def postprocessor_result(self, postprocessor: type | str) -> Any | None:
             postprocessor if isinstance(postprocessor, str) else postprocessor.__name__
         )
         for run in self.postprocessor_runs:
-            if run.name == key:
+            if run.key == key or run.name == key:
                 return run.result
             if not isinstance(postprocessor, str) and issubclass(
                 run.postprocessor_type,
@@ -115,7 +127,8 @@ def postprocessor_results(self, postprocessor: type | str) -> tuple[Any, ...]:
         return tuple(
             run.result
             for run in self.postprocessor_runs
-            if run.name == key
+            if run.key == key
+            or run.name == key
             or (
                 not isinstance(postprocessor, str)
                 and issubclass(run.postprocessor_type, postprocessor)
@@ -143,6 +156,9 @@ class LocalAreaDatasetBuilder:
     variable_cloner: VariableCloner = field(default_factory=VariableCloner)
     postprocessors: tuple[PayloadPostProcessor, ...] = ()
 
+    def __post_init__(self) -> None:
+        _validate_postprocessor_order(self.postprocessors)
+
     def build(
         self,
         *,
@@ -185,6 +201,7 @@ def build(
             result = postprocessor.apply(payload=h5_payload, context=context)
             postprocessor_runs.append(
                 PayloadPostProcessorRun(
+                    key=_postprocessor_spec(postprocessor).key,
                     name=type(postprocessor).__name__,
                     postprocessor_type=type(postprocessor),
                     result=result,
@@ -248,3 +265,31 @@ def _payload_from_postprocessor_result(result: PostProcessorReturn) -> H5Payload
         "Payload postprocessors must return H5Payload or an object exposing "
         "an H5Payload `.payload` attribute"
     )
+
+
+def _validate_postprocessor_order(
+    postprocessors: tuple[PayloadPostProcessor, ...],
+) -> None:
+    seen: set[str] = set()
+    for postprocessor in postprocessors:
+        spec = _postprocessor_spec(postprocessor)
+        if spec.key in seen:
+            raise ValueError(f"Duplicate payload postprocessor key: {spec.key}")
+        missing = tuple(
+            requirement for requirement in spec.requires if requirement not in seen
+        )
+        if missing:
+            raise ValueError(
+                f"{type(postprocessor).__name__} requires postprocessor(s) "
+                f"to run first: {', '.join(missing)}"
+            )
+        seen.add(spec.key)
+
+
+def _postprocessor_spec(
+    postprocessor: PayloadPostProcessor,
+) -> PayloadPostProcessorSpec:
+    spec = getattr(postprocessor, "spec", None)
+    if isinstance(spec, PayloadPostProcessorSpec):
+        return spec
+    return PayloadPostProcessorSpec(key=type(postprocessor).__name__)

policyengine_us_data/build_outputs/us_augmentations.py

@@ -18,13 +18,17 @@
     reported_subsidized_marketplace_by_tax_unit,
 )
 
+from .builder import PayloadPostProcessorSpec
 from .payload import H5Payload, PayloadBuildContext
 from .selection import CloneSelection
 from .simulation_access import calculate_variable_values
 from .variables import GEOGRAPHY_VARIABLES
 
 __all__ = [
     "TAKEUP_VARIABLE_ENTITIES",
+    "US_ENTITY_POSTPROCESSOR_KEY",
+    "US_GEOGRAPHY_POSTPROCESSOR_KEY",
+    "US_TAKEUP_POSTPROCESSOR_KEY",
     "USEntityPostProcessor",
     "USEntityPostProcessorResult",
     "USGeographyPostProcessor",
@@ -42,6 +46,9 @@
     str(spec["variable"]): str(spec["entity"]) for spec in SIMPLE_TAKEUP_VARS
 }
 REQUIRED_TAKEUP_SUBENTITIES = ("tax_unit", "spm_unit")
+US_ENTITY_POSTPROCESSOR_KEY = "us_entity"
+US_GEOGRAPHY_POSTPROCESSOR_KEY = "us_geography"
+US_TAKEUP_POSTPROCESSOR_KEY = "us_takeup"
 
 
 @pipeline_node(
@@ -141,6 +148,8 @@ def data(self) -> PayloadData:
 class USEntityPostProcessor:
     """Apply US entity IDs and calibrated household weights."""
 
+    spec = PayloadPostProcessorSpec(key=US_ENTITY_POSTPROCESSOR_KEY)
+
     def apply(
         self,
         *,
@@ -192,6 +201,8 @@ def apply(
 class USGeographyPostProcessor:
     """Apply block-derived US geography overrides."""
 
+    spec = PayloadPostProcessorSpec(key=US_GEOGRAPHY_POSTPROCESSOR_KEY)
+
     geography_deriver: GeographyDeriver = derive_geography_from_blocks
     _string_geography_variables: tuple[str, ...] = field(
         default=GEOGRAPHY_VARIABLES,
@@ -295,6 +306,10 @@ def _apply_los_angeles_zip_patch(
 class USTakeupPostProcessor:
     """Apply US take-up draws after entity and geography postprocessing."""
 
+    spec = PayloadPostProcessorSpec(
+        key=US_TAKEUP_POSTPROCESSOR_KEY,
+        requires=(US_ENTITY_POSTPROCESSOR_KEY, US_GEOGRAPHY_POSTPROCESSOR_KEY),
+    )
     takeup_applier: TakeupApplier = apply_block_takeup_to_arrays
     sum_person_values_to_tax_units: Callable[
         [np.ndarray, np.ndarray, np.ndarray],
@@ -312,6 +327,7 @@ def apply(
         self._validate_required_subentities(context)
         output = _copy_payload(payload.data)
         time_period = context.time_period
+        self._validate_required_payload_fields(output, time_period)
         results = self._build_takeup_results(output, context)
         takeup_variables = tuple(str(variable) for variable in results)
         self._validate_takeup_variables(takeup_variables)
@@ -349,6 +365,32 @@ def _validate_required_subentities(self, context: PayloadBuildContext) -> None:
                 f"US take-up requires reindexed subentities: {', '.join(missing)}"
             )
 
+    def _validate_required_payload_fields(
+        self,
+        data: PayloadData,
+        time_period: int,
+    ) -> None:
+        _required_period_array(
+            data,
+            "state_fips",
+            time_period,
+            "US take-up requires state_fips from USGeographyPostProcessor",
+        )
+        if _has_period_array(
+            data,
+            "reported_has_subsidized_marketplace_health_coverage_at_interview",
+            time_period,
+        ):
+            for variable in ("person_tax_unit_id", "tax_unit_id"):
+                _required_period_array(
+                    data,
+                    variable,
+                    time_period,
+                    "US take-up reported ACA anchors require "
+                    "person_tax_unit_id and tax_unit_id from "
+                    "USEntityPostProcessor",
+                )
+
     def _build_takeup_results(
         self,
         data: PayloadData,
@@ -481,6 +523,14 @@ def _required_period_array(
     return np.asarray(data[variable][time_period])
 
 
+def _has_period_array(
+    data: Mapping[str, Mapping[Any, np.ndarray]],
+    variable: str,
+    time_period: int,
+) -> bool:
+    return variable in data and time_period in data[variable]
+
+
 def _build_reported_takeup_anchors(
     data: Mapping[str, Mapping[Any, np.ndarray]],
     time_period: int,

tests/unit/build_outputs/test_builder.py

@@ -2,11 +2,13 @@
 from types import SimpleNamespace
 
 import numpy as np
+import pytest
 
 from policyengine_us_data.build_outputs.builder import (
     LocalAreaBuildResult,
     LocalAreaDatasetBuilder,
     PayloadPostProcessorRun,
+    PayloadPostProcessorSpec,
 )
 from policyengine_us_data.build_outputs.payload import H5Payload
 from policyengine_us_data.build_outputs.reindexing import ReindexedEntities
@@ -181,11 +183,13 @@ class _PostProcessor:
         summary={},
         postprocessor_runs=(
             PayloadPostProcessorRun(
+                key="first",
                 name="_PostProcessor",
                 postprocessor_type=_PostProcessor,
                 result=first,
             ),
             PayloadPostProcessorRun(
+                key="second",
                 name="_PostProcessor",
                 postprocessor_type=_PostProcessor,
                 result=second,
@@ -194,9 +198,36 @@ class _PostProcessor:
     )
 
     assert result.postprocessor_result(_PostProcessor) is first
+    assert result.postprocessor_result("second") is second
     assert result.postprocessor_results(_PostProcessor) == (first, second)
 
 
+def test_local_area_dataset_builder_rejects_missing_postprocessor_dependency():
+    class _DependentPostProcessor:
+        spec = PayloadPostProcessorSpec(
+            key="dependent",
+            requires=("upstream",),
+        )
+
+    with pytest.raises(
+        ValueError,
+        match="_DependentPostProcessor requires postprocessor\\(s\\) "
+        "to run first: upstream",
+    ):
+        LocalAreaDatasetBuilder(postprocessors=(_DependentPostProcessor(),))
+
+
+def test_local_area_dataset_builder_rejects_duplicate_postprocessor_keys():
+    class _PostProcessor:
+        spec = PayloadPostProcessorSpec(key="duplicate")
+
+    with pytest.raises(
+        ValueError,
+        match="Duplicate payload postprocessor key: duplicate",
+    ):
+        LocalAreaDatasetBuilder(postprocessors=(_PostProcessor(), _PostProcessor()))
+
+
 def _payload(label: str) -> H5Payload:
     return H5Payload(
         data={label: {2024: np.array([1])}},

tests/unit/build_outputs/test_us_augmentations.py

@@ -13,6 +13,9 @@
     SourceDatasetSnapshot,
 )
 from policyengine_us_data.build_outputs.us_augmentations import (
+    US_ENTITY_POSTPROCESSOR_KEY,
+    US_GEOGRAPHY_POSTPROCESSOR_KEY,
+    US_TAKEUP_POSTPROCESSOR_KEY,
     USEntityPostProcessor,
     USGeographyPostProcessor,
     USTakeupPostProcessor,
@@ -127,11 +130,22 @@ def _geography_payload(context=None) -> H5Payload:
 
 
 def test_default_us_postprocessors_are_in_runtime_order():
-    assert tuple(type(processor) for processor in default_us_postprocessors()) == (
+    postprocessors = default_us_postprocessors()
+
+    assert tuple(type(processor) for processor in postprocessors) == (
         USEntityPostProcessor,
         USGeographyPostProcessor,
         USTakeupPostProcessor,
     )
+    assert tuple(processor.spec.key for processor in postprocessors) == (
+        US_ENTITY_POSTPROCESSOR_KEY,
+        US_GEOGRAPHY_POSTPROCESSOR_KEY,
+        US_TAKEUP_POSTPROCESSOR_KEY,
+    )
+    seen = set()
+    for processor in postprocessors:
+        assert set(processor.spec.requires) <= seen
+        seen.add(processor.spec.key)
 
 
 def test_build_reported_takeup_anchors_skips_missing_period():
@@ -356,6 +370,27 @@ def test_us_takeup_postprocessor_requires_geography_first():
         )
 
 
+def test_us_takeup_postprocessor_requires_entity_ids_for_reported_aca_anchor():
+    payload = _base_payload(
+        {
+            "state_fips": {2024: np.array([6, 37], dtype=np.int32)},
+            "reported_has_subsidized_marketplace_health_coverage_at_interview": {
+                2024: np.array([True, False, False])
+            },
+        }
+    )
+
+    with pytest.raises(
+        ValueError,
+        match="US take-up reported ACA anchors require person_tax_unit_id "
+        "and tax_unit_id from USEntityPostProcessor",
+    ):
+        USTakeupPostProcessor(takeup_applier=lambda **kwargs: {}).apply(
+            payload=payload,
+            context=_context(),
+        )
+
+
 def test_us_takeup_postprocessor_rejects_unknown_takeup_results():
     service = USTakeupPostProcessor(
         takeup_applier=lambda **kwargs: {