Collect Medicaid eligibility from the ACA source batches instead of a second full-population pass

packages/populace-build/tests/test_us_fiscal_refresh_builder.py

@@ -1729,6 +1729,11 @@ def n(self, entity):
             assert entity == "household"
             return 4
 
+        def table(self, entity):
+            # The medicaid eligibility collector indexes the person table.
+            assert entity == "person"
+            return pd.DataFrame({"person_id": np.arange(4, dtype=np.int64)})
+
     monkeypatch.setattr(
         sys,
         "argv",
@@ -1898,7 +1903,9 @@ def n(self, entity):
     monkeypatch.setattr(
         builder,
         "_with_aca_marketplace_source_outputs",
-        lambda frame, specs, *, seed, maximum_microsim_batch_size=None: frame,
+        lambda frame, specs, *, seed, maximum_microsim_batch_size=None, on_batch_simulation=None: (
+            frame
+        ),
     )
     monkeypatch.setattr(
         builder,
@@ -1912,7 +1919,10 @@ def n(self, entity):
     monkeypatch.setattr(
         builder,
         "_with_medicaid_take_up_outputs",
-        lambda frame, specs, *, seed, maximum_microsim_batch_size=None: (frame, {}),
+        lambda frame, specs, *, seed, is_medicaid_eligible=None, maximum_microsim_batch_size=None: (
+            frame,
+            {},
+        ),
     )
     monkeypatch.setattr(
         builder,
@@ -2711,6 +2721,164 @@ def fake_dataset_from_frame(frame_arg, **kwargs):
     assert tax_unit.loc[30, "aca_take_up_rate"] == 0.0
 
 
+def test_aca_source_batches_collect_person_medicaid_eligibility(monkeypatch) -> None:
+    # The Medicaid take-up stage reads person-level is_medicaid_eligible off
+    # the ACA batch simulations instead of running its own full-population
+    # pass; the collector must reassemble per-batch person values in full
+    # person-table order.
+    builder = _load_builder_module()
+    person = pd.DataFrame(
+        {
+            "person_id": np.asarray([1, 2, 3, 4], dtype="int64"),
+            "person_household_id": np.asarray([1, 1, 2, 3], dtype="int64"),
+            "person_tax_unit_id": np.asarray([10, 10, 20, 30], dtype="int64"),
+            "person_spm_unit_id": np.asarray([100, 100, 200, 300], dtype="int64"),
+            "person_family_id": np.asarray([1000, 1000, 2000, 3000], dtype="int64"),
+            "person_marital_unit_id": np.asarray(
+                [10000, 10000, 20000, 30000], dtype="int64"
+            ),
+        }
+    )
+    frame = Frame(
+        {
+            "person": person,
+            "household": pd.DataFrame(
+                {
+                    "household_id": np.asarray([1, 2, 3], dtype="int64"),
+                    "state_fips": np.asarray([1, 1, 2]),
+                }
+            ),
+            "tax_unit": pd.DataFrame(
+                {
+                    "tax_unit_id": np.asarray([10, 20, 30], dtype="int64"),
+                    "stable_tax_unit_draw": [0.1, 0.2, 0.3],
+                }
+            ),
+            "spm_unit": pd.DataFrame({"spm_unit_id": [100, 200, 300]}),
+            "family": pd.DataFrame({"family_id": [1000, 2000, 3000]}),
+            "marital_unit": pd.DataFrame({"marital_unit_id": [10000, 20000, 30000]}),
+        },
+        builder.US_SCHEMA,
+        {
+            "household": builder.Weights(
+                values=np.asarray([10.0, 20.0, 30.0]),
+                kind=WeightKind.DESIGN,
+            )
+        },
+    )
+    target_tables = {
+        builder.US_ACA_APTC_TARGET_TABLE: pd.DataFrame(
+            {
+                "state_fips": ["01"],
+                "target": [3.0],
+            }
+        )
+    }
+    aca_person_eligible = {1: 1.0, 2: 1.0, 3: 1.0, 4: 0.0}
+    medicaid_person_eligible = {1: 0.0, 2: 1.0, 3: 0.0, 4: 1.0}
+
+    class FakeMicrosimulation:
+        def __init__(self, *, dataset):
+            self.dataset = dataset
+
+        def _invalidate_all_caches(self):
+            pass
+
+    def fake_calculate_array(simulation, variable, *, map_to=None):
+        if map_to == "person":
+            source = (
+                medicaid_person_eligible
+                if variable == "is_medicaid_eligible"
+                else aca_person_eligible
+            )
+            return np.asarray(
+                [
+                    source[int(person_id)]
+                    for person_id in simulation.dataset.table("person")["person_id"]
+                ],
+                dtype=np.float64,
+            )
+        assert map_to == "tax_unit"
+        return np.zeros(simulation.dataset.n("tax_unit"), dtype=np.float64)
+
+    monkeypatch.setattr(
+        builder,
+        "_assert_no_formula_owned_columns",
+        lambda frame_arg: None,
+    )
+    monkeypatch.setattr(
+        builder,
+        "_dataset_from_frame",
+        lambda frame_arg, **kwargs: frame_arg,
+    )
+    monkeypatch.setattr(builder, "_calculate_array", fake_calculate_array)
+
+    collector = builder._BatchedPersonEligibilityCollector(
+        frame, "is_medicaid_eligible"
+    )
+    assert collector.values() is None
+    builder._aca_source_tax_unit_table_batched(
+        frame,
+        target_tables,
+        microsimulation_cls=FakeMicrosimulation,
+        maximum_microsim_batch_size=1,
+        on_batch_simulation=collector.collect,
+    )
+    assert collector.values().tolist() == [False, True, False, True]
+
+
+def test_person_eligibility_collector_rejects_partial_coverage() -> None:
+    builder = _load_builder_module()
+    person = pd.DataFrame(
+        {
+            "person_id": np.asarray([1, 2, 3], dtype="int64"),
+            "person_household_id": np.asarray([1, 1, 2], dtype="int64"),
+            "person_tax_unit_id": np.asarray([10, 10, 20], dtype="int64"),
+            "person_spm_unit_id": np.asarray([100, 100, 200], dtype="int64"),
+            "person_family_id": np.asarray([1000, 1000, 2000], dtype="int64"),
+            "person_marital_unit_id": np.asarray([10000, 10000, 20000], dtype="int64"),
+        }
+    )
+    frame = Frame(
+        {
+            "person": person,
+            "household": pd.DataFrame(
+                {
+                    "household_id": np.asarray([1, 2], dtype="int64"),
+                    "state_fips": np.asarray([1, 2]),
+                }
+            ),
+            "tax_unit": pd.DataFrame(
+                {"tax_unit_id": np.asarray([10, 20], dtype="int64")}
+            ),
+            "spm_unit": pd.DataFrame({"spm_unit_id": [100, 200]}),
+            "family": pd.DataFrame({"family_id": [1000, 2000]}),
+            "marital_unit": pd.DataFrame({"marital_unit_id": [10000, 20000]}),
+        },
+        builder.US_SCHEMA,
+        {
+            "household": builder.Weights(
+                values=np.asarray([1.0, 1.0]),
+                kind=WeightKind.DESIGN,
+            )
+        },
+    )
+    collector = builder._BatchedPersonEligibilityCollector(
+        frame, "is_medicaid_eligible"
+    )
+    batch = builder._select_households_by_position(frame, np.asarray([0]))
+    collector.collect(
+        batch,
+        SimpleNamespace(
+            calculate=lambda variable, period=None, map_to=None: np.ones(
+                batch.n("person"), dtype=np.float64
+            )
+        ),
+    )
+    with pytest.raises(RuntimeError, match="covered only 2 of 3 persons"):
+        collector.values()
+
+
 def test_aca_source_runtime_rejects_enrollment_only_fallback() -> None:
     builder = _load_builder_module()
     specs = (
@@ -2784,7 +2952,7 @@ def fake_run_source_stage(
     monkeypatch.setattr(
         builder,
         "_aca_source_tax_unit_table",
-        lambda frame, target_tables, *, simulation=None, maximum_microsim_batch_size=None: (
+        lambda frame, target_tables, *, simulation=None, maximum_microsim_batch_size=None, on_batch_simulation=None: (
             pd.DataFrame({"tax_unit_id": [10, 20], "state_fips": ["06", "06"]})
         ),
     )

packages/populace-build/tests/test_us_medicaid_take_up.py

@@ -584,3 +584,51 @@ def test_person_state_fips_maps_households(self) -> None:
             "36",
             "36",
         ]
+
+    def test_take_up_outputs_use_precomputed_eligibility(self, monkeypatch) -> None:
+        # The release flow collects is_medicaid_eligible during the ACA
+        # source batches; handing it in must skip the dedicated batched
+        # eligibility pass entirely.
+        builder = _load_builder_module()
+        n = 4
+        frame = _frame(n, anchor=np.zeros(n, dtype=bool))
+        household = frame.table("household").copy()
+        household["state_fips"] = [6] * n
+        tables = {entity: frame.table(entity) for entity in frame.entities}
+        tables["household"] = household
+        frame = Frame(
+            tables,
+            frame.schema,
+            {entity: frame.weights_for(entity) for entity in frame.weighted_entities},
+        )
+        specs = (
+            SimpleNamespace(
+                name=(
+                    "cms_medicaid.month2024_12.state_enrollment.06."
+                    "total_medicaid_enrollment@2024"
+                ),
+                value=2 * UNIT_WEIGHT,
+                metadata={
+                    "ledger_geography_level": "state",
+                    "state_fips": "06",
+                    "target_role": "medicaid_enrollment",
+                },
+            ),
+        )
+
+        def fail_eligibility_pass(*args, **kwargs):
+            raise AssertionError(
+                "precomputed eligibility must not trigger the batched pass"
+            )
+
+        monkeypatch.setattr(
+            builder, "_medicaid_person_eligibility", fail_eligibility_pass
+        )
+        result, diagnostics = builder._with_medicaid_take_up_outputs(
+            frame,
+            specs,
+            seed=7,
+            is_medicaid_eligible=np.asarray([True, True, True, False]),
+        )
+        assert US_MEDICAID_TAKE_UP_VARIABLE in result.table("person")
+        assert diagnostics["national"]["eligible_weight"] == 3 * UNIT_WEIGHT