Allocate Medicaid cost if enrolled from SLCSP (#1138)

* Allocate Medicaid conditional cost from SLCSP

* Preserve Medicaid costs in local H5 builds

* Pin us-data to PE-US Medicaid cost fix

Author: MaxGhenis

changelog.d/1138.changed.md

@@ -0,0 +1 @@
+Allocate conditional Medicaid cost from an SLCSP age/location index in us-data.

policyengine_us_data/build_outputs/us_augmentations.py

@@ -31,11 +31,14 @@
     "TAKEUP_VARIABLE_ENTITIES",
     "US_ENTITY_POSTPROCESSOR_KEY",
     "US_GEOGRAPHY_POSTPROCESSOR_KEY",
+    "US_MEDICAID_COST_POSTPROCESSOR_KEY",
     "US_TAKEUP_POSTPROCESSOR_KEY",
     "USEntityPostProcessor",
     "USEntityPostProcessorResult",
     "USGeographyPostProcessor",
     "USGeographyPostProcessorResult",
+    "USMedicaidCostPostProcessor",
+    "USMedicaidCostPostProcessorResult",
     "USTakeupPostProcessor",
     "USTakeupPostProcessorResult",
     "default_us_postprocessors",
@@ -52,6 +55,7 @@
 US_ENTITY_POSTPROCESSOR_KEY = "us_entity"
 US_GEOGRAPHY_POSTPROCESSOR_KEY = "us_geography"
 US_TAKEUP_POSTPROCESSOR_KEY = "us_takeup"
+US_MEDICAID_COST_POSTPROCESSOR_KEY = "us_medicaid_cost"
 
 
 @pipeline_node(
@@ -134,6 +138,32 @@ def data(self) -> PayloadData:
         return self.payload.data
 
 
+@pipeline_node(
+    id="local_h5_us_medicaid_cost_postprocessor_result",
+    label="USMedicaidCostPostProcessorResult",
+    node_type="library",
+    description="US Medicaid conditional-cost local H5 payload data.",
+    source_file="policyengine_us_data/build_outputs/us_augmentations.py",
+    status="current",
+    stability="moving",
+    pathways=["local_h5"],
+    validation_commands=[
+        "uv run pytest tests/unit/build_outputs/test_us_augmentations.py"
+    ],
+)
+@dataclass(frozen=True)
+class USMedicaidCostPostProcessorResult:
+    """Payload after conditional Medicaid cost fields are applied."""
+
+    payload: H5Payload
+
+    @property
+    def data(self) -> PayloadData:
+        """Augmented payload data retained for transitional callers."""
+
+        return self.payload.data
+
+
 @pipeline_node(
     id="local_h5_us_entity_postprocessor",
     label="USEntityPostProcessor",
@@ -571,22 +601,106 @@ def _build_eligibility_masks(
         }
 
 
+@pipeline_node(
+    id="local_h5_us_medicaid_cost_postprocessor",
+    label="USMedicaidCostPostProcessor",
+    node_type="library",
+    description="Preserve Medicaid cost-if-enrolled inputs in local H5 payloads.",
+    source_file="policyengine_us_data/build_outputs/us_augmentations.py",
+    status="current",
+    stability="moving",
+    pathways=["local_h5"],
+    validation_commands=[
+        "uv run pytest tests/unit/build_outputs/test_us_augmentations.py"
+    ],
+)
+@dataclass(frozen=True)
+class USMedicaidCostPostProcessor:
+    """Preserve source Medicaid conditional costs after local H5 transforms."""
+
+    spec = PayloadPostProcessorSpec(
+        key=US_MEDICAID_COST_POSTPROCESSOR_KEY,
+        requires=(
+            US_ENTITY_POSTPROCESSOR_KEY,
+            US_GEOGRAPHY_POSTPROCESSOR_KEY,
+            US_TAKEUP_POSTPROCESSOR_KEY,
+        ),
+    )
+
+    def apply(
+        self,
+        *,
+        payload: H5Payload,
+        context: PayloadBuildContext,
+    ) -> USMedicaidCostPostProcessorResult:
+        """Return a payload with source conditional Medicaid costs preserved."""
+
+        output = _copy_payload(payload.data)
+        cost_periods = output.get("medicaid_cost_if_enrolled", {})
+        if context.time_period not in cost_periods:
+            source_values = _source_person_period_values(
+                context=context,
+                variable="medicaid_cost_if_enrolled",
+            )
+            if source_values is not None:
+                output["medicaid_cost_if_enrolled"] = {
+                    context.time_period: source_values.astype(np.float32)
+                }
+
+        variable_entities = dict(payload.variable_entities)
+        if "medicaid_cost_if_enrolled" in output:
+            variable_entities["medicaid_cost_if_enrolled"] = "person"
+        return USMedicaidCostPostProcessorResult(
+            payload=H5Payload(
+                data=output,
+                time_period=payload.time_period,
+                entity_lengths=payload.entity_lengths,
+                variable_entities=variable_entities,
+            ),
+        )
+
+
 def default_us_postprocessors() -> tuple[
-    USEntityPostProcessor | USGeographyPostProcessor | USTakeupPostProcessor, ...
+    USEntityPostProcessor
+    | USGeographyPostProcessor
+    | USTakeupPostProcessor
+    | USMedicaidCostPostProcessor,
+    ...,
 ]:
     """Return production US postprocessors in their required order."""
 
     return (
         USEntityPostProcessor(),
         USGeographyPostProcessor(),
         USTakeupPostProcessor(),
+        USMedicaidCostPostProcessor(),
     )
 
 
 def _copy_payload(data: Mapping[str, Mapping[Any, np.ndarray]]) -> PayloadData:
     return {variable: dict(periods) for variable, periods in data.items()}
 
 
+def _source_person_period_values(
+    *,
+    context: PayloadBuildContext,
+    variable: str,
+) -> np.ndarray | None:
+    """Return a source person input mapped to the local payload, if available."""
+
+    if variable not in context.source.input_variables:
+        return None
+    provider = context.source.variable_provider
+    get_array = getattr(provider, "get_array", None)
+    if not callable(get_array):
+        return None
+    try:
+        values = np.asarray(get_array(variable, context.time_period))
+    except (KeyError, ValueError):
+        return None
+    return values[context.reindexed.person_source_indices]
+
+
 def _required_period_array(
     data: Mapping[str, Mapping[Any, np.ndarray]],
     variable: str,

policyengine_us_data/datasets/cps/cps.py

@@ -49,6 +49,9 @@
     derive_treasury_tipped_occupation_code,
     derive_is_tipped_occupation,
 )
+from policyengine_us_data.datasets.cps.medicaid_cost import (
+    add_medicaid_cost_if_enrolled_to_dataset,
+)
 from policyengine_us_data.utils.takeup import (
     _sum_person_values_to_tax_units,
     _voluntary_filing_age_bin,
@@ -360,6 +363,8 @@ def generate(self):
         add_takeup(self)
         logging.info("Imputing Marketplace plan benchmark ratio")
         add_marketplace_plan_benchmark_ratio(self)
+        logging.info("Adding Medicaid cost if enrolled")
+        add_medicaid_cost_if_enrolled_to_dataset(self)
         logging.info("Deriving other health insurance premiums")
         derive_other_health_insurance_premiums(self)
         logging.info("Downsampling")

policyengine_us_data/datasets/cps/extended_cps.py

@@ -25,6 +25,9 @@
     derive_flsa_overtime_premium,
     load_take_up_rate,
 )
+from policyengine_us_data.datasets.cps.medicaid_cost import (
+    add_medicaid_cost_if_enrolled_to_time_period_data,
+)
 from policyengine_us_data.datasets.cps.takeup import prioritize_reported_recipients
 from policyengine_us_data.datasets.org import (
     ORG_IMPUTED_VARIABLES,
@@ -1190,6 +1193,10 @@ def generate(self):
             new_data,
             self.time_period,
         )
+        new_data = add_medicaid_cost_if_enrolled_to_time_period_data(
+            new_data,
+            self.time_period,
+        )
         new_data = self._validate_housing_assistance_microsimulation(
             new_data,
             self.time_period,