Add Medicare Part B enrollment count target

changelog.d/716.fixed

@@ -0,0 +1 @@
+Adds a CMS-sourced Medicare Part B enrollment count target to complement the beneficiary-paid premium calibration target.

policyengine_us_data/calibration/target_config.yaml

@@ -93,6 +93,9 @@ include:
     geo_level: national
   - variable: medicaid
     geo_level: national
+  - variable: person_count
+    geo_level: national
+    domain_variable: medicare_enrolled
   - variable: medicare_part_b_premiums
     geo_level: national
   - variable: other_medical_expenses

policyengine_us_data/db/etl_national_targets.py

@@ -16,6 +16,9 @@
     get_beneficiary_paid_medicare_part_b_premiums_notes,
     get_beneficiary_paid_medicare_part_b_premiums_source,
     get_beneficiary_paid_medicare_part_b_premiums_target,
+    get_medicare_part_b_enrollment_notes,
+    get_medicare_part_b_enrollment_source,
+    get_medicare_part_b_enrollment_target,
 )
 from policyengine_us_data.utils.db import (
     DEFAULT_YEAR,
@@ -330,6 +333,13 @@ def extract_national_targets(year: int = DEFAULT_YEAR):
             "notes": "ACA Premium Tax Credit recipients",
             "year": HARDCODED_YEAR,
         },
+        {
+            "constraint_variable": "medicare_enrolled",
+            "person_count": get_medicare_part_b_enrollment_target(HARDCODED_YEAR),
+            "source": get_medicare_part_b_enrollment_source(HARDCODED_YEAR),
+            "notes": get_medicare_part_b_enrollment_notes(HARDCODED_YEAR),
+            "year": HARDCODED_YEAR,
+        },
         {
             "constraint_variable": "spm_unit_energy_subsidy_reported",
             "target_variable": "household_count",

policyengine_us_data/utils/cms_medicare.py

@@ -2,6 +2,10 @@
     2024: 139.837e9,
 }
 
+MEDICARE_PART_B_ENROLLMENT_TARGETS = {
+    2024: 62_084_000,
+}
+
 
 MEDICARE_STATE_BUY_IN_MINIMUM_BENEFICIARIES = {
     2024: 10_000_000,
@@ -13,6 +17,32 @@
 }
 
 
+def get_medicare_part_b_enrollment_target(year: int) -> float:
+    try:
+        return MEDICARE_PART_B_ENROLLMENT_TARGETS[year]
+    except KeyError as exc:
+        raise ValueError(
+            f"No Medicare Part B enrollment target sourced for {year}."
+        ) from exc
+
+
+def get_medicare_part_b_enrollment_source(year: int) -> str:
+    enrollment = MEDICARE_PART_B_ENROLLMENT_TARGETS[year]
+    return (
+        "CMS 2024 Medicare Trustees Report Table V.B3 intermediate estimate "
+        f"for {year} Part B enrollment ({enrollment:,.0f} beneficiaries)"
+    )
+
+
+def get_medicare_part_b_enrollment_notes(year: int) -> str:
+    return (
+        "Total Medicare Part B enrollment count. This is intentionally a "
+        "separate calibration anchor from beneficiary-paid Part B premiums, "
+        "because some enrollees have premiums paid on their behalf through "
+        "MSP/state buy-in pathways."
+    )
+
+
 def get_beneficiary_paid_medicare_part_b_premiums_target(year: int) -> float:
     try:
         return BENEFICIARY_PAID_MEDICARE_PART_B_PREMIUM_TARGETS[year]

policyengine_us_data/utils/loss.py

@@ -14,6 +14,7 @@
 )
 from policyengine_us_data.utils.cms_medicare import (
     get_beneficiary_paid_medicare_part_b_premiums_target,
+    get_medicare_part_b_enrollment_target,
 )
 from policyengine_us_data.db.etl_irs_soi import get_national_geography_soi_target
 from policyengine_core.reforms import Reform
@@ -110,6 +111,15 @@
     ],
 }
 
+
+def _add_medicare_part_b_enrollment_target(loss_matrix, targets_array, sim):
+    label = "nation/cms/medicare_part_b_enrollment"
+    enrolled = sim.calculate("medicare_enrolled", map_to="person", period=2025).values
+    loss_matrix[label] = sim.map_result(enrolled.astype(float), "person", "household")
+    targets_array.append(get_medicare_part_b_enrollment_target(2024))
+    return targets_array, loss_matrix
+
+
 ACA_SPENDING_TARGETS = {
     2024: 98e9,
 }
@@ -574,6 +584,12 @@ def build_loss_matrix(dataset: type, time_period):
 
     targets_array.append(aca_enrollment_target)
 
+    targets_array, loss_matrix = _add_medicare_part_b_enrollment_target(
+        loss_matrix,
+        targets_array,
+        sim,
+    )
+
     # Treasury EITC
 
     loss_matrix["nation/treasury/eitc"] = sim.calculate(

tests/unit/calibration/test_loss_targets.py

@@ -3,6 +3,7 @@
 import pytest
 
 from policyengine_us_data.utils.loss import (
+    _add_medicare_part_b_enrollment_target,
     _get_aca_national_targets,
     _add_ctc_targets,
     _get_medicaid_national_targets,
@@ -89,6 +90,26 @@ def map_result(self, values, source_entity, target_entity, how=None):
         return np.asarray(values, dtype=np.float32)
 
 
+class _FakeMedicareEnrollmentSimulation:
+    def __init__(self):
+        self.calculate_calls = []
+        self.map_result_calls = []
+
+    def calculate(self, variable, map_to=None, period=None):
+        self.calculate_calls.append((variable, map_to, period))
+        if variable != "medicare_enrolled":
+            raise AssertionError(f"Unexpected variable {variable!r}")
+        if map_to != "person":
+            raise AssertionError(f"Unexpected map_to {map_to!r}")
+        return _FakeArrayResult([1.0, 0.0, 1.0])
+
+    def map_result(self, values, source_entity, target_entity, how=None):
+        self.map_result_calls.append((source_entity, target_entity, how))
+        assert source_entity == "person"
+        assert target_entity == "household"
+        return np.asarray(values, dtype=np.float32)
+
+
 def test_add_ctc_targets(monkeypatch):
     monkeypatch.setattr(
         "policyengine_us_data.utils.loss.get_national_geography_soi_target",
@@ -123,3 +144,23 @@ def test_add_ctc_targets(monkeypatch):
         loss_matrix["nation/irs/non_refundable_ctc_count"],
         np.array([1.0, 1.0, 0.0], dtype=np.float32),
     )
+
+
+def test_add_medicare_part_b_enrollment_target(monkeypatch):
+    monkeypatch.setattr(
+        "policyengine_us_data.utils.loss.get_medicare_part_b_enrollment_target",
+        lambda year: 62_084_000.0,
+    )
+    sim = _FakeMedicareEnrollmentSimulation()
+
+    targets, loss_matrix = _add_medicare_part_b_enrollment_target(
+        pd.DataFrame(),
+        [],
+        sim,
+    )
+
+    assert targets == [62_084_000.0]
+    np.testing.assert_array_equal(
+        loss_matrix["nation/cms/medicare_part_b_enrollment"],
+        np.array([1.0, 0.0, 1.0], dtype=np.float32),
+    )

tests/unit/test_cms_medicare_targets.py

@@ -4,6 +4,9 @@
     get_beneficiary_paid_medicare_part_b_premiums_notes,
     get_beneficiary_paid_medicare_part_b_premiums_source,
     get_beneficiary_paid_medicare_part_b_premiums_target,
+    get_medicare_part_b_enrollment_notes,
+    get_medicare_part_b_enrollment_source,
+    get_medicare_part_b_enrollment_target,
 )
 
 
@@ -23,3 +26,19 @@ def test_beneficiary_paid_medicare_part_b_notes_describe_out_of_pocket_semantics
     notes = get_beneficiary_paid_medicare_part_b_premiums_notes(2024)
     assert "out-of-pocket" in notes
     assert "gross trust-fund premium income" in notes
+
+
+def test_medicare_part_b_enrollment_target_2024_is_sourced():
+    assert get_medicare_part_b_enrollment_target(2024) == pytest.approx(62_084_000)
+
+
+def test_medicare_part_b_enrollment_source_mentions_trustees_table():
+    source = get_medicare_part_b_enrollment_source(2024)
+    assert "2024 Medicare Trustees Report" in source
+    assert "Table V.B3" in source
+
+
+def test_medicare_part_b_enrollment_notes_describe_complementary_semantics():
+    notes = get_medicare_part_b_enrollment_notes(2024)
+    assert "separate calibration anchor" in notes
+    assert "MSP/state buy-in" in notes

tests/unit/test_etl_national_targets.py

@@ -199,3 +199,66 @@ def test_load_national_targets_supports_liheap_household_counts(tmp_path, monkey
         ).first()
         assert liheap_target is not None
         assert liheap_target.value == 5_876_646
+
+
+def test_load_national_targets_supports_medicare_part_b_enrollment_counts(
+    tmp_path, monkeypatch
+):
+    calibration_dir = tmp_path / "calibration"
+    calibration_dir.mkdir()
+    db_uri = f"sqlite:///{calibration_dir / 'policy_data.db'}"
+    engine = create_database(db_uri)
+
+    with Session(engine) as session:
+        national = _make_stratum(session, notes="United States")
+        assert national is not None
+
+    monkeypatch.setattr(
+        "policyengine_us_data.db.etl_national_targets.STORAGE_FOLDER",
+        tmp_path,
+    )
+
+    conditional_targets = [
+        {
+            "constraint_variable": "medicare_enrolled",
+            "person_count": 62_084_000,
+            "source": "CMS 2024 Medicare Trustees Report Table V.B3",
+            "notes": "Total Medicare Part B enrollment count",
+            "year": 2024,
+        }
+    ]
+
+    load_national_targets(
+        direct_targets_df=pd.DataFrame(),
+        tax_filer_df=pd.DataFrame(),
+        tax_expenditure_df=pd.DataFrame(),
+        conditional_targets=conditional_targets,
+    )
+
+    with Session(engine) as session:
+        medicare_stratum = session.exec(
+            select(Stratum).where(
+                Stratum.notes == "National medicare_enrolled Recipients"
+            )
+        ).first()
+        assert medicare_stratum is not None
+
+        constraints = {
+            (
+                constraint.constraint_variable,
+                constraint.operation,
+                constraint.value,
+            )
+            for constraint in medicare_stratum.constraints_rel
+        }
+        assert ("medicare_enrolled", ">", "0") in constraints
+
+        medicare_target = session.exec(
+            select(Target).where(
+                Target.stratum_id == medicare_stratum.stratum_id,
+                Target.variable == "person_count",
+                Target.period == 2024,
+            )
+        ).first()
+        assert medicare_target is not None
+        assert medicare_target.value == 62_084_000