fix cms_quality_measures to stage diagnosis codes from core.condition

core.medical_claim has no diagnosis columns — normalized_code and
principal_diagnosis_code live in core.condition. Added two staging steps
that enrich claims before measure functions consume them:

- _stg_medical_claim: LEFT JOIN rank-1 condition for principal_diagnosis_code
- _stg_medical_claim_condition: INNER JOIN all conditions for MCC cohort matching

Updated uamcc_int_mcc_cohort, uamcc_int_planned_admission,
uamcc_int_outcome_exclusion, and uamcc_int_numerator to consume
the staging tables instead of raw core.medical_claim.

src/aco/pipe/cms_quality_measures.py

@@ -0,0 +1,529 @@
+from __future__ import annotations
+
+from aco.express import cms_quality_measures as ex
+from aco.express.base import Expr
+from aco.pipe.base import Pipeline
+from aco.table.cms_quality_measures import (
+    CmsQmAcrIntIndexAdmission,
+    CmsQmAcrIntPlannedReadmission,
+    CmsQmAcrIntSpecialtyCohort,
+    CmsQmAcrPerformancePeriod,
+    CmsQmAcrSummary,
+    CmsQmHwrIntDenominator,
+    CmsQmHwrIntPlannedReadmission,
+    CmsQmHwrPerformancePeriod,
+    CmsQmHwrSummary,
+    CmsQmStgMedicalClaim,
+    CmsQmStgMedicalClaimCondition,
+    CmsQmUamccIntDenominator,
+    CmsQmUamccIntDenominatorExclusion,
+    CmsQmUamccIntMccCohort,
+    CmsQmUamccIntNumerator,
+    CmsQmUamccIntOutcomeExclusion,
+    CmsQmUamccIntPersonTime,
+    CmsQmUamccIntPlannedAdmission,
+    CmsQmUamccPerformancePeriod,
+    CmsQmUamccSummary,
+)
+from bib.tag import Tag
+
+# ── Shared citation refs ─────────────────────────────────────────────────────
+
+_REACH_REFS = [
+    Tag.module("aco"),
+    Tag.program("reach"),
+    Tag.source("cms-website"),
+]
+
+_MIPS_REFS = [
+    Tag.module("aco"),
+    Tag.program("mips"),
+    Tag.source("cms-website"),
+]
+
+_UAMCC_REFS = [
+    *_REACH_REFS,
+    Tag.measure("UAMCC"),
+]
+
+_ACR_REFS = [
+    *_REACH_REFS,
+    Tag.measure("ACR"),
+]
+
+_HWR_REFS = [
+    *_MIPS_REFS,
+    Tag.measure("HWR"),
+]
+
+# ── UAMCC pipeline ───────────────────────────────────────────────────────────
+# All-Cause Unplanned Admissions for Patients with Multiple Chronic Conditions
+# NQF #2888 — ACO REACH + MIPS MCC
+#
+# Measure pipeline order:
+#   1. performance_period         → anchor dates + lookback window
+#   2. int_mcc_cohort             → 9 chronic condition group membership
+#   3. int_denominator            → MCC-eligible beneficiaries (≥66, 2+ groups)
+#   4. int_denominator_exclusion  → hospice, enrollment, ACO visit exclusions
+#   5. int_planned_admission      → PAA v4.0 2024 classification
+#   6. int_outcome_exclusion      → planned + injury + complication flags
+#   7. int_person_time            → at-risk days / person-years
+#   8. int_numerator              → qualifying unplanned admissions
+#   9. summary                    → observed rate per 100 person-years
+
+# ── Staging pipelines ────────────────────────────────────────────────────────
+# core.medical_claim has no diagnosis codes — those live in core.condition.
+# These staging steps enrich claims before the measure pipelines consume them.
+
+_STG_STEPS = [
+    Expr(
+        name="cms_quality_measures._stg_medical_claim",
+        fn=ex.stg_medical_claim,
+        output=CmsQmStgMedicalClaim,
+        after=[
+            "core.medical_claim",
+            "core.condition",
+        ],
+        refs=_REACH_REFS,
+        description=(
+            "Medical claims enriched with the principal (rank-1) diagnosis "
+            "code from core.condition.  Preserves the claim-line grain of "
+            "core.medical_claim."
+        ),
+    ),
+    Expr(
+        name="cms_quality_measures._stg_medical_claim_condition",
+        fn=ex.stg_medical_claim_condition,
+        output=CmsQmStgMedicalClaimCondition,
+        after=[
+            "core.medical_claim",
+            "core.condition",
+        ],
+        refs=_REACH_REFS,
+        description=(
+            "One row per (claim, diagnosis code) pair — core.medical_claim "
+            "joined to all core.condition rows.  Used for MCC chronic-condition "
+            "cohort matching where every diagnosis position must be checked."
+        ),
+    ),
+]
+
+_UAMCC_VALUE_SETS = [
+    "cms_quality_measures._uamcc_value_set_cohort",
+    "cms_quality_measures._uamcc_value_set_exclusions",
+    "cms_quality_measures._uamcc_value_set_paa1",
+    "cms_quality_measures._uamcc_value_set_paa2",
+    "cms_quality_measures._uamcc_value_set_paa3",
+    "cms_quality_measures._uamcc_value_set_paa4",
+    "cms_quality_measures._uamcc_value_set_ccs_icd10_cm",
+    "cms_quality_measures._uamcc_value_set_ccs_icd10_pcs",
+]
+
+pipeline_uamcc = Pipeline(
+    exprs=[
+        *_STG_STEPS,
+        Expr(
+            name="cms_quality_measures._uamcc_performance_period",
+            fn=ex.uamcc_performance_period,
+            output=CmsQmUamccPerformancePeriod,
+            after=["cms_quality_measures._uamcc_performance_period"],
+            refs=_UAMCC_REFS,
+            description=(
+                "Single-row anchor for the UAMCC 12-month measurement window. "
+                "Carries measure_id (UAMCC), NQF #2888, performance_year, "
+                "performance_period_begin/end, and the prior-year lookback "
+                "window for chronic condition identification."
+            ),
+        ),
+        Expr(
+            name="cms_quality_measures._uamcc_int_mcc_cohort",
+            fn=ex.uamcc_int_mcc_cohort,
+            output=CmsQmUamccIntMccCohort,
+            after=[
+                "cms_quality_measures._stg_medical_claim_condition",
+                *_UAMCC_VALUE_SETS[:1],  # _uamcc_value_set_cohort
+            ],
+            refs=_UAMCC_REFS,
+            description=(
+                "Identifies each beneficiary's qualifying chronic condition groups "
+                "by joining medical claims to the UAMCC Cohort value set (ICD-10 codes "
+                "for the nine MCC groups: AMI, Alzheimer's, AFib, CKD, COPD/asthma, "
+                "Depression, Diabetes, Heart failure, Stroke/TIA). "
+                "Emits one row per (person_id, chronic_condition_group) pair with "
+                "earliest qualifying claim date and claim count."
+            ),
+        ),
+        Expr(
+            name="cms_quality_measures._uamcc_int_denominator",
+            fn=ex.uamcc_int_denominator,
+            output=CmsQmUamccIntDenominator,
+            after=[
+                "cms_quality_measures._uamcc_int_mcc_cohort",
+                "core.patient",
+                "cms_quality_measures._uamcc_performance_period",
+            ],
+            refs=_UAMCC_REFS,
+            description=(
+                "Builds the UAMCC denominator: Medicare FFS beneficiaries aged ≥66 at "
+                "the start of the measurement period who have diagnoses in two or more "
+                "of the nine MCC chronic disease groups. "
+                "Filters cohort to ≥2 distinct condition groups and joins to patient "
+                "age at the performance period start date."
+            ),
+        ),
+        Expr(
+            name="cms_quality_measures._uamcc_int_denominator_exclusion",
+            fn=ex.uamcc_int_denominator_exclusion,
+            output=CmsQmUamccIntDenominatorExclusion,
+            after=[
+                "cms_quality_measures._uamcc_int_denominator",
+                "core.patient",
+            ],
+            refs=_UAMCC_REFS,
+            description=(
+                "Identifies beneficiaries excluded from the UAMCC denominator per "
+                "MIF §3.10: voluntary alignment after Jan 1, missing Part A/B "
+                "enrollment, hospice enrollment during prior year or at period start, "
+                "no qualifying E&M visit with the aligned ACO in measurement or prior "
+                "year, and beneficiaries with no at-risk person-time during the year."
+            ),
+        ),
+        Expr(
+            name="cms_quality_measures._uamcc_int_planned_admission",
+            fn=ex.uamcc_int_planned_admission,
+            output=CmsQmUamccIntPlannedAdmission,
+            after=[
+                "cms_quality_measures._stg_medical_claim",
+                *_UAMCC_VALUE_SETS[2:8],  # paa1-4 + ccs maps
+            ],
+            refs=_UAMCC_REFS,
+            description=(
+                "Applies Planned Admission Algorithm (PAA) v4.0 2024 to classify "
+                "inpatient admissions as planned or unplanned. "
+                "Rule 1: always-planned procedure (PAA1, 5 CCS categories). "
+                "Rule 2: always-planned principal diagnosis (PAA2, 4 CCS categories). "
+                "Rule 3: potentially-planned procedure (PAA3, ~2,705 entries) AND "
+                "the principal diagnosis is NOT an acute diagnosis (PAA4, ~11,369 entries)."
+            ),
+        ),
+        Expr(
+            name="cms_quality_measures._uamcc_int_outcome_exclusion",
+            fn=ex.uamcc_int_outcome_exclusion,
+            output=CmsQmUamccIntOutcomeExclusion,
+            after=[
+                "cms_quality_measures._stg_medical_claim",
+                "cms_quality_measures._uamcc_int_planned_admission",
+                "cms_quality_measures._uamcc_value_set_exclusions",
+                "cms_quality_measures._uamcc_value_set_ccs_icd10_cm",
+            ],
+            refs=_UAMCC_REFS,
+            description=(
+                "Flags inpatient admissions excluded from the UAMCC numerator per "
+                "MIF §3.7: planned (PAA), direct-from-SNF/rehab, 10-day buffer period, "
+                "hospice, procedure/surgery complications (CCS 145, 237, 238, 257), "
+                "accidents/injuries (CCS E-codes 2601–2621), and admissions before "
+                "the beneficiary's first qualifying visit with the aligned ACO."
+            ),
+        ),
+        Expr(
+            name="cms_quality_measures._uamcc_int_person_time",
+            fn=ex.uamcc_int_person_time,
+            output=CmsQmUamccIntPersonTime,
+            after=[
+                "cms_quality_measures._uamcc_int_denominator",
+                "core.encounter",
+                "cms_quality_measures._uamcc_performance_period",
+            ],
+            refs=_UAMCC_REFS,
+            description=(
+                "Calculates at-risk person-time (days and person-years) for each "
+                "eligible beneficiary. Excludes time spent in hospital, SNF, or "
+                "acute rehabilitation; the 10-day post-discharge buffer; and time "
+                "after hospice enrollment. "
+                "Person-years = at_risk_days / 365.25. "
+                "The UAMCC outcome is a rate per 100 person-years."
+            ),
+        ),
+        Expr(
+            name="cms_quality_measures._uamcc_int_numerator",
+            fn=ex.uamcc_int_numerator,
+            output=CmsQmUamccIntNumerator,
+            after=[
+                "cms_quality_measures._stg_medical_claim",
+                "cms_quality_measures._uamcc_int_denominator",
+                "cms_quality_measures._uamcc_int_outcome_exclusion",
+                "cms_quality_measures._uamcc_value_set_ccs_icd10_cm",
+            ],
+            refs=_UAMCC_REFS,
+            description=(
+                "Identifies qualifying unplanned acute inpatient admissions for the "
+                "UAMCC numerator: denominator-eligible beneficiaries whose admission "
+                "is not in the outcome exclusion set. "
+                "Each qualifying admission counts as one event; the final rate is "
+                "expressed per 100 person-years."
+            ),
+        ),
+        Expr(
+            name="cms_quality_measures.uamcc_summary",
+            fn=ex.uamcc_summary,
+            output=CmsQmUamccSummary,
+            after=[
+                "cms_quality_measures._uamcc_int_numerator",
+                "cms_quality_measures._uamcc_int_person_time",
+                "cms_quality_measures._uamcc_int_denominator",
+                "cms_quality_measures._uamcc_performance_period",
+            ],
+            refs=_UAMCC_REFS,
+            description=(
+                "Aggregates UAMCC observed admission rate per 100 person-years. "
+                "Computes denominator count, total person-years, observed admission "
+                "count, and crude observed rate. "
+                "The full risk-standardized acute admission rate (RSAAR, NQF #2888) "
+                "requires CMS's hierarchical negative-binomial model fit across all "
+                "ACO data; expected_admissions and rsaar are NULL placeholders."
+            ),
+        ),
+    ]
+)
+
+# ── ACR pipeline ─────────────────────────────────────────────────────────────
+# Risk-Standardized, All-Condition Readmission — NQF #1789 — ACO REACH
+#
+# Measure pipeline order:
+#   1. performance_period        → anchor dates
+#   2. int_index_admission       → eligible hospitalizations (denominator)
+#   3. int_specialty_cohort      → 5-cohort CCS assignment
+#   4. int_planned_readmission   → PAA v4.0 classification
+#   5. summary                   → observed readmission rate
+
+_ACR_VALUE_SETS = [
+    "cms_quality_measures._acr_value_set_cohort_ccs",
+    "cms_quality_measures._acr_value_set_cohort_icd10",
+    "cms_quality_measures._acr_value_set_exclusions",
+    "cms_quality_measures._acr_value_set_paa1",
+    "cms_quality_measures._acr_value_set_paa2",
+    "cms_quality_measures._acr_value_set_paa3",
+    "cms_quality_measures._acr_value_set_paa4",
+    "cms_quality_measures._acr_value_set_ccs_icd10_cm",
+    "cms_quality_measures._acr_value_set_ccs_icd10_pcs",
+]
+
+pipeline_acr = Pipeline(
+    exprs=[
+        Expr(
+            name="cms_quality_measures._acr_performance_period",
+            fn=ex.acr_performance_period,
+            output=CmsQmAcrPerformancePeriod,
+            after=["cms_quality_measures._acr_performance_period"],
+            refs=_ACR_REFS,
+            description=(
+                "Single-row anchor for the ACR 12-month measurement window. "
+                "Carries measure_id (ACR), NQF #1789, performance_year, and "
+                "performance_period_begin/end. "
+                "Quarterly rolling 12-month rates are also calculated for "
+                "informational reporting."
+            ),
+        ),
+        Expr(
+            name="cms_quality_measures._acr_int_index_admission",
+            fn=ex.acr_int_index_admission,
+            output=CmsQmAcrIntIndexAdmission,
+            after=[
+                "core.encounter",
+                "cms_quality_measures._acr_value_set_exclusions",
+                "cms_quality_measures._acr_value_set_ccs_icd10_cm",
+            ],
+            refs=_ACR_REFS,
+            description=(
+                "Identifies eligible index hospitalizations for the ACR denominator. "
+                "Includes Medicare FFS beneficiaries ≥65 discharged alive from "
+                "non-federal, short-stay acute-care or critical access hospitals "
+                "who were not transferred to another acute care facility. "
+                "Applies cohort-level CCS exclusions (~49 categories) from the "
+                "ACR value set. A readmission may also qualify as an index admission."
+            ),
+        ),
+        Expr(
+            name="cms_quality_measures._acr_int_specialty_cohort",
+            fn=ex.acr_int_specialty_cohort,
+            output=CmsQmAcrIntSpecialtyCohort,
+            after=[
+                "cms_quality_measures._acr_int_index_admission",
+                "cms_quality_measures._acr_value_set_cohort_ccs",
+                "cms_quality_measures._acr_value_set_cohort_icd10",
+                "core.procedure",
+            ],
+            refs=_ACR_REFS,
+            description=(
+                "Assigns each ACR index admission to one of five mutually exclusive "
+                "specialty cohorts using AHRQ CCS. "
+                "Priority: Surgery/Gynecology (ICD-10-PCS, ~1,683 codes) always wins; "
+                "then Cardiorespiratory, Cardiovascular, Neurology by principal diagnosis "
+                "CCS (~279 CCS entries); Medicine is the default. "
+                "Cohorts reflect how care for patients is organized within hospitals."
+            ),
+        ),
+        Expr(
+            name="cms_quality_measures._acr_int_planned_readmission",
+            fn=ex.acr_int_planned_readmission,
+            output=CmsQmAcrIntPlannedReadmission,
+            after=[
+                "core.encounter",
+                "cms_quality_measures._acr_int_index_admission",
+                "cms_quality_measures._acr_value_set_paa1",
+                "cms_quality_measures._acr_value_set_paa2",
+                "cms_quality_measures._acr_value_set_paa3",
+                "cms_quality_measures._acr_value_set_paa4",
+                "cms_quality_measures._acr_value_set_ccs_icd10_cm",
+                "cms_quality_measures._acr_value_set_ccs_icd10_pcs",
+            ],
+            refs=_ACR_REFS,
+            description=(
+                "Applies PAA v4.0 to all hospitalizations within 30 days of an "
+                "index discharge to classify readmissions as planned or unplanned. "
+                "Rule 1: always-planned procedure (PA1). "
+                "Rule 2: always-planned principal diagnosis (PA2). "
+                "Rule 3: potentially-planned procedure (PA3, ~2,701 entries) AND NOT "
+                "acute principal diagnosis (PA4, ~11,369 entries). "
+                "Readmissions to psychiatric or rehabilitation facilities are always "
+                "excluded from the ACR outcome."
+            ),
+        ),
+        Expr(
+            name="cms_quality_measures.acr_summary",
+            fn=ex.acr_summary,
+            output=CmsQmAcrSummary,
+            after=[
+                "cms_quality_measures._acr_int_index_admission",
+                "cms_quality_measures._acr_int_planned_readmission",
+                "cms_quality_measures._acr_performance_period",
+            ],
+            refs=_ACR_REFS,
+            description=(
+                "Aggregates ACR observed unplanned 30-day readmission rate for the ACO. "
+                "Computes denominator (eligible index admissions), numerator (unplanned "
+                "readmissions), and crude observed rate. "
+                "The full risk-standardized readmission rate (RSRR, NQF #1789) requires "
+                "CMS's hierarchical logistic model; expected_readmissions and rsrr are "
+                "NULL placeholders."
+            ),
+        ),
+    ]
+)
+
+# ── HWR pipeline ─────────────────────────────────────────────────────────────
+# Hospital-wide, 30-Day, All-cause Unplanned Readmission — MIPS Groups
+#
+# Measure pipeline order:
+#   1. performance_period         → MIPS year anchor
+#   2. int_denominator            → eligible hospitalizations + specialty cohort
+#   3. int_planned_readmission    → PAA v4.0 classification
+#   4. summary                    → observed readmission rate by TIN
+
+_HWR_VALUE_SETS = [
+    "cms_quality_measures._hwr_value_set_specialty_cohort",
+    "cms_quality_measures._hwr_value_set_surg_gyn_cohort",
+    "cms_quality_measures._hwr_value_set_cohort_exclusions",
+    "cms_quality_measures._hwr_value_set_paa1",
+    "cms_quality_measures._hwr_value_set_paa2",
+    "cms_quality_measures._hwr_value_set_paa3",
+    "cms_quality_measures._hwr_value_set_paa4",
+]
+
+pipeline_hwr = Pipeline(
+    exprs=[
+        Expr(
+            name="cms_quality_measures._hwr_performance_period",
+            fn=ex.hwr_performance_period,
+            output=CmsQmHwrPerformancePeriod,
+            after=["cms_quality_measures._hwr_performance_period"],
+            refs=_HWR_REFS,
+            description=(
+                "Single-row anchor for the MIPS HWR performance year. "
+                "Covers Jan 1 – Dec 31 of the MIPS performance year. "
+                "The HWR measure attributes readmissions to MIPS clinician groups "
+                "via NPIs and TINs."
+            ),
+        ),
+        Expr(
+            name="cms_quality_measures._hwr_int_denominator",
+            fn=ex.hwr_int_denominator,
+            output=CmsQmHwrIntDenominator,
+            after=[
+                "core.encounter",
+                "cms_quality_measures._hwr_value_set_cohort_exclusions",
+                "cms_quality_measures._hwr_value_set_specialty_cohort",
+                "cms_quality_measures._hwr_value_set_surg_gyn_cohort",
+                "core.procedure",
+            ],
+            refs=_HWR_REFS,
+            description=(
+                "Builds the MIPS HWR denominator: Medicare FFS beneficiaries ≥65 "
+                "discharged alive from non-federal, short-stay acute-care or critical "
+                "access hospitals during the performance period who were not transferred "
+                "to another acute care facility. "
+                "Applies cohort exclusions (~231 CCS categories) and assigns each "
+                "admission to one of five specialty cohorts using AHRQ CCS. "
+                "Attribution TIN fields are populated downstream by CMS attribution logic."
+            ),
+        ),
+        Expr(
+            name="cms_quality_measures._hwr_int_planned_readmission",
+            fn=ex.hwr_int_planned_readmission,
+            output=CmsQmHwrIntPlannedReadmission,
+            after=[
+                "core.encounter",
+                "cms_quality_measures._hwr_int_denominator",
+                "cms_quality_measures._hwr_value_set_paa1",
+                "cms_quality_measures._hwr_value_set_paa2",
+                "cms_quality_measures._hwr_value_set_paa3",
+                "cms_quality_measures._hwr_value_set_paa4",
+                "cms_quality_measures._acr_value_set_ccs_icd10_cm",
+            ],
+            refs=_HWR_REFS,
+            description=(
+                "Applies PAA v4.0 to hospitalizations within 30 days of an HWR index "
+                "discharge. The MIPS HWR PAA is equivalent to the ACR PAA and shares "
+                "the same four-rule algorithm. "
+                "Uses HWR-specific value sets (PR.1–PR.4) maintained separately by CMS. "
+                "Readmissions to psychiatric or rehabilitation facilities are always "
+                "excluded from the HWR outcome."
+            ),
+        ),
+        Expr(
+            name="cms_quality_measures.hwr_summary",
+            fn=ex.hwr_summary,
+            output=CmsQmHwrSummary,
+            after=[
+                "cms_quality_measures._hwr_int_denominator",
+                "cms_quality_measures._hwr_int_planned_readmission",
+                "cms_quality_measures._hwr_performance_period",
+            ],
+            refs=_HWR_REFS,
+            description=(
+                "Aggregates MIPS HWR observed unplanned 30-day readmission rate. "
+                "The measure attributes readmissions to three clinician group types: "
+                "Discharge Clinician Group (discharge procedure within last 3 days), "
+                "Primary Inpatient Care Provider (most charges during stay), and "
+                "Outpatient PCP (most primary care E&M visits in prior 12 months). "
+                "The full RSRR requires the hierarchical logistic model; "
+                "expected_readmissions and rsrr are NULL placeholders."
+            ),
+        ),
+    ]
+)
+
+# Combined pipeline for running all three measures together
+# Staging steps are already in pipeline_uamcc.exprs
+pipeline = Pipeline(
+    exprs=[
+        *pipeline_uamcc.exprs,
+        *pipeline_acr.exprs,
+        *pipeline_hwr.exprs,
+    ]
+)
+
+run = pipeline.run
+run_uamcc = pipeline_uamcc.run
+run_acr = pipeline_acr.run
+run_hwr = pipeline_hwr.run

src/aco/table/cms_quality_measures.py

@@ -38,6 +38,47 @@ from decimal import Decimal
 
 from aco.table.base import SQLTable
 
+# ── Staging ──────────────────────────────────────────────────────────────────
+
+
+class CmsQmStgMedicalClaim(SQLTable):
+    """Schema: cms_quality_measures / Table: _stg_medical_claim
+
+    Medical claims enriched with the principal (rank-1) diagnosis code from
+    ``core.condition``.  Preserves the claim-line grain of
+    ``core.medical_claim``; every line of the same claim carries the same
+    ``principal_diagnosis_code``.
+    """
+
+    __schema__ = "cms_quality_measures"
+    __tablename__ = "_stg_medical_claim"
+
+    claim_id: str | None = None
+    person_id: str | None = None
+    claim_start_date: date | None = None
+    claim_end_date: date | None = None
+    hcpcs_code: str | None = None
+    place_of_service_code: str | None = None
+    principal_diagnosis_code: str | None = None
+
+
+class CmsQmStgMedicalClaimCondition(SQLTable):
+    """Schema: cms_quality_measures / Table: _stg_medical_claim_condition
+
+    One row per (claim, diagnosis code) pair — ``core.medical_claim`` joined
+    to all ``core.condition`` rows.  Used for MCC chronic-condition cohort
+    matching where every diagnosis position must be checked.
+    """
+
+    __schema__ = "cms_quality_measures"
+    __tablename__ = "_stg_medical_claim_condition"
+
+    claim_id: str | None = None
+    person_id: str | None = None
+    claim_start_date: date | None = None
+    normalized_code: str | None = None
+
+
 # ── UAMCC Performance Period ─────────────────────────────────────────────────
 
 

tests/aco/test_express_cms_quality_measures.py

@@ -23,6 +23,8 @@ from aco.express.cms_quality_measures import (
     hwr_int_denominator,
     hwr_int_planned_readmission,
     hwr_performance_period,
+    stg_medical_claim,
+    stg_medical_claim_condition,
     uamcc_int_denominator,
     uamcc_int_denominator_exclusion,
     uamcc_int_mcc_cohort,
@@ -164,8 +166,8 @@ def uamcc_cohort_value_set_df() -> pl.DataFrame:
 
 
 @pytest.fixture
-def medical_claim_df() -> pl.DataFrame:
+def stg_medical_claim_df() -> pl.DataFrame:
-    """Medical claims referencing chronic condition and procedure codes."""
+    """Staged medical claims with principal diagnosis (rank-1 condition)."""
     return pl.DataFrame(
         {
             "claim_id": [
@@ -208,16 +210,6 @@ def medical_claim_df() -> pl.DataFrame:
                 date(2024, 3, 4),
                 date(2024, 7, 3),
             ],
-            "normalized_code": [
-                "I21.0",  # P001 AMI
-                "I50.1",  # P001 Heart failure → 2 conditions
-                "E11.9",  # P002 Diabetes
-                "N18.3",  # P002 CKD → 2 conditions
-                "J44.1",  # P004 COPD
-                "I48.0",  # P004 AFib → 2 conditions
-                "G30.0",  # P005 Alzheimer
-                "F32.0",  # P005 Depression → 2 conditions
-            ],
             "principal_diagnosis_code": [
                 "I21.0",
                 "I50.1",
@@ -234,6 +226,55 @@ def medical_claim_df() -> pl.DataFrame:
     )
 
 
+@pytest.fixture
+def stg_medical_claim_condition_df() -> pl.DataFrame:
+    """Staged claim-condition pairs (all diagnosis positions)."""
+    return pl.DataFrame(
+        {
+            "claim_id": [
+                "CLM001",
+                "CLM002",
+                "CLM003",
+                "CLM004",
+                "CLM005",
+                "CLM006",
+                "CLM007",
+                "CLM008",
+            ],
+            "person_id": [
+                "P001",
+                "P001",
+                "P002",
+                "P002",
+                "P004",
+                "P004",
+                "P005",
+                "P005",
+            ],
+            "claim_start_date": [
+                date(2024, 3, 1),
+                date(2024, 5, 1),
+                date(2024, 2, 1),
+                date(2024, 4, 1),
+                date(2024, 1, 1),
+                date(2024, 6, 1),
+                date(2024, 3, 1),
+                date(2024, 7, 1),
+            ],
+            "normalized_code": [
+                "I21.0",  # P001 AMI
+                "I50.1",  # P001 Heart failure → 2 conditions
+                "E11.9",  # P002 Diabetes
+                "N18.3",  # P002 CKD → 2 conditions
+                "J44.1",  # P004 COPD
+                "I48.0",  # P004 AFib → 2 conditions
+                "G30.0",  # P005 Alzheimer
+                "F32.0",  # P005 Depression → 2 conditions
+            ],
+        }
+    )
+
+
 @pytest.fixture
 def encounter_df() -> pl.DataFrame:
     """Core encounters including acute inpatient stays."""
@@ -461,6 +502,133 @@ class TestModuleConstants:
         assert 2605 in _INJURY_ACCIDENT_CCS  # Firearm
 
 
+# ═══════════════════════════════════════════════════════════════════════════════
+# Staging tests
+# ═══════════════════════════════════════════════════════════════════════════════
+
+
+class TestStgMedicalClaim:
+    """Tests for stg_medical_claim staging function."""
+
+    def test_adds_principal_diagnosis_code(self) -> None:
+        medical_claim = pl.DataFrame(
+            {
+                "claim_id": ["CLM001", "CLM002"],
+                "person_id": ["P001", "P002"],
+                "claim_start_date": [date(2024, 1, 1), date(2024, 2, 1)],
+                "claim_end_date": [date(2024, 1, 3), date(2024, 2, 3)],
+                "hcpcs_code": pl.Series([None, None], dtype=pl.String),
+                "place_of_service_code": ["21", "21"],
+            }
+        )
+        condition = pl.DataFrame(
+            {
+                "claim_id": ["CLM001", "CLM001", "CLM002"],
+                "normalized_code": ["I21.0", "I50.1", "E11.9"],
+                "condition_rank": [1, 2, 1],
+            }
+        )
+        result = stg_medical_claim(medical_claim, condition)
+        assert "principal_diagnosis_code" in result.columns
+        row1 = result.filter(pl.col("claim_id") == "CLM001")
+        assert row1["principal_diagnosis_code"][0] == "I21.0"
+
+    def test_preserves_claim_line_grain(self) -> None:
+        medical_claim = pl.DataFrame(
+            {
+                "claim_id": ["CLM001", "CLM001"],
+                "person_id": ["P001", "P001"],
+                "claim_start_date": [date(2024, 1, 1), date(2024, 1, 1)],
+                "claim_end_date": [date(2024, 1, 3), date(2024, 1, 3)],
+                "hcpcs_code": ["99213", "71046"],
+                "place_of_service_code": ["21", "21"],
+            }
+        )
+        condition = pl.DataFrame(
+            {
+                "claim_id": ["CLM001"],
+                "normalized_code": ["I21.0"],
+                "condition_rank": [1],
+            }
+        )
+        result = stg_medical_claim(medical_claim, condition)
+        assert len(result) == 2
+
+    def test_null_when_no_condition(self) -> None:
+        medical_claim = pl.DataFrame(
+            {
+                "claim_id": ["CLM001"],
+                "person_id": ["P001"],
+                "claim_start_date": [date(2024, 1, 1)],
+                "claim_end_date": [date(2024, 1, 3)],
+                "hcpcs_code": pl.Series([None], dtype=pl.String),
+                "place_of_service_code": ["21"],
+            }
+        )
+        condition = pl.DataFrame(
+            {
+                "claim_id": pl.Series([], dtype=pl.String),
+                "normalized_code": pl.Series([], dtype=pl.String),
+                "condition_rank": pl.Series([], dtype=pl.Int64),
+            }
+        )
+        result = stg_medical_claim(medical_claim, condition)
+        assert len(result) == 1
+        assert result["principal_diagnosis_code"][0] is None
+
+
+class TestStgMedicalClaimCondition:
+    """Tests for stg_medical_claim_condition staging function."""
+
+    def test_fans_out_to_all_conditions(self) -> None:
+        medical_claim = pl.DataFrame(
+            {
+                "claim_id": ["CLM001", "CLM001"],
+                "person_id": ["P001", "P001"],
+                "claim_start_date": [date(2024, 1, 1), date(2024, 1, 1)],
+                "claim_end_date": [date(2024, 1, 3), date(2024, 1, 3)],
+                "hcpcs_code": ["99213", "71046"],
+                "place_of_service_code": ["21", "21"],
+            }
+        )
+        condition = pl.DataFrame(
+            {
+                "claim_id": ["CLM001", "CLM001"],
+                "normalized_code": ["I21.0", "I50.1"],
+                "condition_rank": [1, 2],
+            }
+        )
+        result = stg_medical_claim_condition(medical_claim, condition)
+        assert len(result) == 2
+        assert set(result["normalized_code"].to_list()) == {"I21.0", "I50.1"}
+
+    def test_expected_columns(self) -> None:
+        medical_claim = pl.DataFrame(
+            {
+                "claim_id": ["CLM001"],
+                "person_id": ["P001"],
+                "claim_start_date": [date(2024, 1, 1)],
+                "claim_end_date": [date(2024, 1, 3)],
+                "hcpcs_code": pl.Series([None], dtype=pl.String),
+                "place_of_service_code": ["21"],
+            }
+        )
+        condition = pl.DataFrame(
+            {
+                "claim_id": ["CLM001"],
+                "normalized_code": ["I21.0"],
+                "condition_rank": [1],
+            }
+        )
+        result = stg_medical_claim_condition(medical_claim, condition)
+        assert set(result.columns) == {
+            "claim_id",
+            "person_id",
+            "claim_start_date",
+            "normalized_code",
+        }
+
+
 # ═══════════════════════════════════════════════════════════════════════════════
 # UAMCC tests
 # ═══════════════════════════════════════════════════════════════════════════════
@@ -514,15 +682,19 @@ class TestUamccIntMccCohort:
     """Tests for uamcc_int_mcc_cohort chronic condition identification."""
 
     def test_returns_dataframe(
-        self, medical_claim_df, uamcc_cohort_value_set_df
+        self, stg_medical_claim_condition_df, uamcc_cohort_value_set_df
     ) -> None:
-        result = uamcc_int_mcc_cohort(medical_claim_df, uamcc_cohort_value_set_df)
+        result = uamcc_int_mcc_cohort(
+            stg_medical_claim_condition_df, uamcc_cohort_value_set_df
+        )
         assert isinstance(result, pl.DataFrame)
 
     def test_expected_columns(
-        self, medical_claim_df, uamcc_cohort_value_set_df
+        self, stg_medical_claim_condition_df, uamcc_cohort_value_set_df
     ) -> None:
-        result = uamcc_int_mcc_cohort(medical_claim_df, uamcc_cohort_value_set_df)
+        result = uamcc_int_mcc_cohort(
+            stg_medical_claim_condition_df, uamcc_cohort_value_set_df
+        )
         for col in [
             "person_id",
             "chronic_condition_group",
@@ -534,9 +706,11 @@ class TestUamccIntMccCohort:
             assert col in result.columns
 
     def test_identifies_ami_for_p001(
-        self, medical_claim_df, uamcc_cohort_value_set_df
+        self, stg_medical_claim_condition_df, uamcc_cohort_value_set_df
     ) -> None:
-        result = uamcc_int_mcc_cohort(medical_claim_df, uamcc_cohort_value_set_df)
+        result = uamcc_int_mcc_cohort(
+            stg_medical_claim_condition_df, uamcc_cohort_value_set_df
+        )
         p001 = result.filter(
             (pl.col("person_id") == "P001")
             & (pl.col("chronic_condition_group") == "AMI")
@@ -544,9 +718,11 @@ class TestUamccIntMccCohort:
         assert len(p001) == 1
 
     def test_identifies_heart_failure_for_p001(
-        self, medical_claim_df, uamcc_cohort_value_set_df
+        self, stg_medical_claim_condition_df, uamcc_cohort_value_set_df
     ) -> None:
-        result = uamcc_int_mcc_cohort(medical_claim_df, uamcc_cohort_value_set_df)
+        result = uamcc_int_mcc_cohort(
+            stg_medical_claim_condition_df, uamcc_cohort_value_set_df
+        )
         p001_hf = result.filter(
             (pl.col("person_id") == "P001")
             & (pl.col("chronic_condition_group") == "HEART_FAILURE")
@@ -554,9 +730,11 @@ class TestUamccIntMccCohort:
         assert len(p001_hf) == 1
 
     def test_p001_has_two_conditions(
-        self, medical_claim_df, uamcc_cohort_value_set_df
+        self, stg_medical_claim_condition_df, uamcc_cohort_value_set_df
     ) -> None:
-        result = uamcc_int_mcc_cohort(medical_claim_df, uamcc_cohort_value_set_df)
+        result = uamcc_int_mcc_cohort(
+            stg_medical_claim_condition_df, uamcc_cohort_value_set_df
+        )
         p001 = result.filter(pl.col("person_id") == "P001")
         assert len(p001) == 2
 
@@ -566,11 +744,7 @@ class TestUamccIntMccCohort:
                 "claim_id": ["CLM999"],
                 "person_id": ["P999"],
                 "claim_start_date": [date(2024, 1, 1)],
-                "claim_end_date": [date(2024, 1, 1)],
                 "normalized_code": ["Z99.99"],
-                "principal_diagnosis_code": ["Z99.99"],
-                "hcpcs_code": pl.Series([None], dtype=pl.String),
-                "place_of_service_code": ["11"],
             }
         )
         result = uamcc_int_mcc_cohort(no_match, uamcc_cohort_value_set_df)
@@ -583,11 +757,7 @@ class TestUamccIntMccCohort:
                 "claim_id": ["CLM_A", "CLM_B"],
                 "person_id": ["P001", "P001"],
                 "claim_start_date": [date(2024, 1, 1), date(2024, 3, 1)],
-                "claim_end_date": [date(2024, 1, 2), date(2024, 3, 2)],
                 "normalized_code": ["I21.0", "I21.1"],
-                "principal_diagnosis_code": ["I21.0", "I21.1"],
-                "hcpcs_code": pl.Series([None, None], dtype=pl.String),
-                "place_of_service_code": ["21", "21"],
             }
         )
         result = uamcc_int_mcc_cohort(two_ami_claims, uamcc_cohort_value_set_df)
@@ -760,9 +930,9 @@ class TestUamccIntDenominatorExclusion:
 class TestUamccIntPlannedAdmission:
     """Tests for uamcc_int_planned_admission PAA algorithm."""
 
-    def test_returns_dataframe(self, medical_claim_df) -> None:
+    def test_returns_dataframe(self, stg_medical_claim_df) -> None:
         result = uamcc_int_planned_admission(
-            medical_claim_df,
+            stg_medical_claim_df,
             _make_paa1_df(),
             _make_paa2_df(),
             _make_paa3_df(),
@@ -772,9 +942,9 @@ class TestUamccIntPlannedAdmission:
         )
         assert isinstance(result, pl.DataFrame)
 
-    def test_expected_columns(self, medical_claim_df) -> None:
+    def test_expected_columns(self, stg_medical_claim_df) -> None:
         result = uamcc_int_planned_admission(
-            medical_claim_df,
+            stg_medical_claim_df,
             _make_paa1_df(),
             _make_paa2_df(),
             _make_paa3_df(),
@@ -799,7 +969,6 @@ class TestUamccIntPlannedAdmission:
                 "person_id": ["P001"],
                 "claim_start_date": [date(2025, 3, 1)],
                 "claim_end_date": [date(2025, 3, 3)],
-                "normalized_code": ["Z51.11"],
                 "principal_diagnosis_code": ["Z51.11"],
                 "hcpcs_code": pl.Series([None], dtype=pl.String),
                 "place_of_service_code": ["21"],
@@ -818,10 +987,10 @@ class TestUamccIntPlannedAdmission:
         assert row["is_planned"][0] == 1
         assert row["planned_rule"][0] == "RULE2"
 
-    def test_unplanned_acute_mi(self, medical_claim_df) -> None:
+    def test_unplanned_acute_mi(self, stg_medical_claim_df) -> None:
         """AMI claim (I21.0) is an acute diagnosis → unplanned."""
         result = uamcc_int_planned_admission(
-            medical_claim_df,
+            stg_medical_claim_df,
             _make_paa1_df(),
             _make_paa2_df(),
             _make_paa3_df(),
@@ -833,9 +1002,9 @@ class TestUamccIntPlannedAdmission:
         assert ami_claim["is_planned"][0] == 0
         assert ami_claim["planned_rule"][0] is None
 
-    def test_unplanned_row_count_matches_input(self, medical_claim_df) -> None:
+    def test_unplanned_row_count_matches_input(self, stg_medical_claim_df) -> None:
         result = uamcc_int_planned_admission(
-            medical_claim_df,
+            stg_medical_claim_df,
             _make_paa1_df(),
             _make_paa2_df(),
             _make_paa3_df(),
@@ -843,7 +1012,7 @@ class TestUamccIntPlannedAdmission:
             _make_ccs_icd10_cm_df(),
             _make_ccs_icd10_pcs_df(),
         )
-        assert len(result) == len(medical_claim_df)
+        assert len(result) == len(stg_medical_claim_df)
 
 
 class TestUamccIntOutcomeExclusion:
@@ -858,7 +1027,6 @@ class TestUamccIntOutcomeExclusion:
                 "person_id": ["P001"],
                 "claim_start_date": [date(2025, 4, 1)],
                 "claim_end_date": [date(2025, 4, 3)],
-                "normalized_code": ["T82.7XXA"],
                 "principal_diagnosis_code": ["T82.7XXA"],
                 "hcpcs_code": pl.Series([None], dtype=pl.String),
                 "place_of_service_code": ["21"],
@@ -874,7 +1042,6 @@ class TestUamccIntOutcomeExclusion:
                 "person_id": ["P001"],
                 "claim_start_date": [date(2025, 5, 1)],
                 "claim_end_date": [date(2025, 5, 2)],
-                "normalized_code": ["W19.XXXA"],
                 "principal_diagnosis_code": ["W19.XXXA"],
                 "hcpcs_code": pl.Series([None], dtype=pl.String),
                 "place_of_service_code": ["21"],
@@ -937,10 +1104,10 @@ class TestUamccIntOutcomeExclusion:
         assert result["is_injury_or_accident"][0] == 1
 
     def test_non_excluded_claim_not_returned(
-        self, medical_claim_df, exclusions_value_set_df, empty_planned_df
+        self, stg_medical_claim_df, exclusions_value_set_df, empty_planned_df
     ) -> None:
         # Heart failure (CCS 108) — not in exclusion list
-        hf_only = medical_claim_df.filter(pl.col("claim_id") == "CLM002")
+        hf_only = stg_medical_claim_df.filter(pl.col("claim_id") == "CLM002")
         result = uamcc_int_outcome_exclusion(
             hf_only,
             empty_planned_df,
@@ -950,7 +1117,7 @@ class TestUamccIntOutcomeExclusion:
         assert len(result) == 0
 
     def test_planned_claim_flagged(
-        self, medical_claim_df, exclusions_value_set_df
+        self, stg_medical_claim_df, exclusions_value_set_df
     ) -> None:
         planned_df = pl.DataFrame(
             {
@@ -959,7 +1126,7 @@ class TestUamccIntOutcomeExclusion:
             }
         )
         result = uamcc_int_outcome_exclusion(
-            medical_claim_df.filter(pl.col("claim_id") == "CLM001"),
+            stg_medical_claim_df.filter(pl.col("claim_id") == "CLM001"),
             planned_df,
             exclusions_value_set_df,
             _make_ccs_icd10_cm_df(),