stack/notebooks/acodb_explorer.py

import marimo

__generated_with = "0.21.1"
app = marimo.App(width="full")


@app.cell(hide_code=True)
def _():
    import marimo as mo

    return (mo,)


@app.cell(hide_code=True)
def _(mo):
    mo.md("""
    # ACO Data Explorer

    Exploring a **1,000-patient Medicare CCLF** dataset processed through the
    [Tuva Health](https://thetuvaproject.com/) dbt analytics framework.

    The database contains **1,007 tables** across **22 schemas** covering claims preprocessing,
    chronic conditions, quality measures, risk adjustment, financial PMPM, readmissions, and more.
    """)
    return


@app.cell(hide_code=True)
def _():
    import altair as alt
    import polars as pl

    from conf import connect

    connect.theme()
    _s3 = connect.obstore()  # noqa: F841 — marimo auto-discovers in Files

    from fhirworx import PALETTE  # noqa: E402 — needs theme() to add styles to path

    con = connect.duckdb()

    def q(sql):
        """Run a query and return a Polars DataFrame."""
        return con.execute(sql).pl()

    return PALETTE, alt, pl, q


@app.cell(hide_code=True)
def _(mo):
    mo.md("""
    ## Database Overview
    """)
    return


@app.cell(hide_code=True)
def _(mo, q):
    schema_counts = q("""
        SELECT table_schema AS schema, count(*) AS tables
        FROM information_schema.tables
        WHERE table_schema != 'information_schema'
        GROUP BY table_schema
        ORDER BY tables DESC
    """)
    mo.ui.table(schema_counts, label="Tables per Schema")
    return


@app.cell(hide_code=True)
def _(mo, q):
    core_sizes = q("""
        SELECT 'patient' AS table_name, count(*) AS rows FROM core.patient
        UNION ALL SELECT 'encounter', count(*) FROM core.encounter
        UNION ALL SELECT 'condition', count(*) FROM core.condition
        UNION ALL SELECT 'medical_claim', count(*) FROM core.medical_claim
        UNION ALL SELECT 'pharmacy_claim', count(*) FROM core.pharmacy_claim
        UNION ALL SELECT 'procedure', count(*) FROM core.procedure
        UNION ALL SELECT 'lab_result', count(*) FROM core.lab_result
        UNION ALL SELECT 'observation', count(*) FROM core.observation
        UNION ALL SELECT 'practitioner', count(*) FROM core.practitioner
        UNION ALL SELECT 'eligibility', count(*) FROM core.eligibility
        UNION ALL SELECT 'member_months', count(*) FROM core.member_months
        ORDER BY rows DESC
    """)
    mo.ui.table(core_sizes, label="Core Table Row Counts")
    return


@app.cell(hide_code=True)
def _(mo):
    mo.md("""
    ## Patient Demographics
    """)
    return


@app.cell(hide_code=True)
def _(PALETTE, alt, mo, q):
    patients = q(
        "SELECT sex, race, state, age, age_group FROM core.patient WHERE state IS NOT NULL"
    )

    sex_chart = (
        alt.Chart(patients.to_pandas())
        .mark_arc(innerRadius=50)
        .encode(
            theta=alt.Theta("count():Q"),
            color=alt.Color("sex:N", scale=alt.Scale(range=PALETTE)),
            tooltip=["sex:N", "count():Q"],
        )
        .properties(title="Sex Distribution", width=250, height=250)
    )

    race_chart = (
        alt.Chart(patients.to_pandas())
        .mark_bar()
        .encode(
            x=alt.X("count():Q", title="Patients"),
            y=alt.Y("race:N", sort="-x", title=None),
            color=alt.Color("race:N", scale=alt.Scale(range=PALETTE), legend=None),
            tooltip=["race:N", "count():Q"],
        )
        .properties(title="Race Distribution", width=350, height=250)
    )

    mo.hstack([sex_chart, race_chart], justify="start", gap=2)
    return (patients,)


@app.cell(hide_code=True)
def _(PALETTE, alt, mo, patients):
    age_chart = (
        alt.Chart(patients.to_pandas())
        .mark_bar()
        .encode(
            x=alt.X("age:Q", bin=alt.Bin(step=5), title="Age"),
            y=alt.Y("count():Q", title="Patients"),
            color=alt.Color("sex:N", scale=alt.Scale(range=PALETTE)),
            tooltip=["count():Q"],
        )
        .properties(title="Age Distribution by Sex", width=600, height=300)
    )

    state_chart = (
        alt.Chart(
            patients.group_by("state")
            .len()
            .sort("len", descending=True)
            .head(15)
            .to_pandas()
        )
        .mark_bar()
        .encode(
            x=alt.X("len:Q", title="Patients"),
            y=alt.Y("state:N", sort="-x", title=None),
            color=alt.value(PALETTE[2]),
            tooltip=["state:N", "len:Q"],
        )
        .properties(title="Top 15 States", width=350, height=300)
    )

    mo.hstack([age_chart, state_chart], justify="start", gap=2)
    return


@app.cell(hide_code=True)
def _(mo):
    mo.md("""
    ## Encounter Analysis
    """)
    return


@app.cell(hide_code=True)
def _(PALETTE, alt, mo, q):
    enc_types = q("""
        SELECT encounter_type, count(*) AS encounters
        FROM core.encounter
        GROUP BY encounter_type
        ORDER BY encounters DESC
    """)

    enc_chart = (
        alt.Chart(enc_types.to_pandas())
        .mark_bar()
        .encode(
            x=alt.X("encounters:Q", title="Count"),
            y=alt.Y("encounter_type:N", sort="-x", title=None),
            color=alt.value(PALETTE[7]),
            tooltip=["encounter_type:N", "encounters:Q"],
        )
        .properties(title="Encounters by Type", width=600, height=450)
    )

    enc_chart
    return


@app.cell(hide_code=True)
def _(PALETTE, alt, mo, q):
    enc_monthly = q("""
        SELECT strftime(encounter_start_date, '%Y-%m') AS month,
               encounter_type,
               count(*) AS encounters
        FROM core.encounter
        WHERE encounter_type IN (
            'office visit', 'emergency department', 'acute inpatient',
            'outpatient hospital or clinic', 'outpatient surgery'
        )
        GROUP BY month, encounter_type
        ORDER BY month
    """)

    trend_chart = (
        alt.Chart(enc_monthly.to_pandas())
        .mark_line(point=True)
        .encode(
            x=alt.X("month:T", title="Month"),
            y=alt.Y("encounters:Q", title="Encounters"),
            color=alt.Color(
                "encounter_type:N", title="Type", scale=alt.Scale(range=PALETTE)
            ),
            tooltip=["month:T", "encounter_type:N", "encounters:Q"],
        )
        .properties(title="Monthly Encounter Trends (Key Types)", width=800, height=350)
    )

    trend_chart
    return


@app.cell(hide_code=True)
def _(mo):
    mo.md("""
    ## Chronic Conditions
    """)
    return


@app.cell(hide_code=True)
def _(PALETTE, alt, mo, q):
    chronic = q("""
        SELECT condition, count(DISTINCT person_id) AS patients
        FROM chronic_conditions.cms_chronic_conditions_long
        GROUP BY condition
        ORDER BY patients DESC
        LIMIT 20
    """)

    chronic_chart = (
        alt.Chart(chronic.to_pandas())
        .mark_bar()
        .encode(
            x=alt.X("patients:Q", title="Patients"),
            y=alt.Y("condition:N", sort="-x", title=None),
            color=alt.Color(
                "patients:Q",
                scale=alt.Scale(range=[PALETTE[1] + "33", PALETTE[1]]),
                legend=None,
            ),
            tooltip=["condition:N", "patients:Q"],
        )
        .properties(
            title="Top 20 Chronic Conditions by Patient Count", width=700, height=500
        )
    )

    chronic_chart
    return


@app.cell(hide_code=True)
def _(mo, q):
    comorbidity = q("""
        SELECT person_id, count(DISTINCT condition) AS condition_count
        FROM chronic_conditions.cms_chronic_conditions_long
        GROUP BY person_id
    """)

    mo.md(f"""
    ### Comorbidity Burden

    | Metric | Value |
    |--------|-------|
    | Patients with chronic conditions | **{comorbidity.height}** |
    | Mean conditions per patient | **{comorbidity["condition_count"].mean():.1f}** |
    | Median conditions per patient | **{comorbidity["condition_count"].median():.0f}** |
    | Max conditions per patient | **{comorbidity["condition_count"].max()}** |
    """)
    return (comorbidity,)


@app.cell(hide_code=True)
def _(PALETTE, alt, comorbidity):
    burden_chart = (
        alt.Chart(comorbidity.to_pandas())
        .mark_bar()
        .encode(
            x=alt.X(
                "condition_count:Q",
                bin=alt.Bin(maxbins=20),
                title="Number of Chronic Conditions",
            ),
            y=alt.Y("count():Q", title="Patients"),
            color=alt.value(PALETTE[5]),
            tooltip=["count():Q"],
        )
        .properties(title="Distribution of Comorbidity Burden", width=600, height=300)
    )

    burden_chart
    return


@app.cell(hide_code=True)
def _(mo):
    mo.md("""
    ## CMS-HCC Risk Scores
    """)
    return


@app.cell(hide_code=True)
def _(PALETTE, alt, mo, pl, q):
    risk = q("""
        SELECT r.person_id, r.payment_risk_score, r.v24_risk_score,
               r.normalized_risk_score, r.member_months,
               p.sex, p.age_group, p.age
        FROM cms_hcc.patient_risk_scores r
        JOIN core.patient p ON r.person_id = p.person_id
        WHERE r.payment_risk_score IS NOT NULL
    """)

    # Convert Decimal columns to float for Altair compatibility
    risk = risk.with_columns(
        [
            pl.col("payment_risk_score").cast(pl.Float64),
            pl.col("v24_risk_score").cast(pl.Float64),
            pl.col("normalized_risk_score").cast(pl.Float64),
            pl.col("member_months").cast(pl.Float64),
        ]
    )

    risk_hist = (
        alt.Chart(risk.to_pandas())
        .mark_bar()
        .encode(
            x=alt.X(
                "payment_risk_score:Q",
                bin=alt.Bin(maxbins=30),
                title="Payment Risk Score",
            ),
            y=alt.Y("count():Q", title="Patients"),
            color=alt.Color("sex:N", scale=alt.Scale(range=PALETTE)),
            tooltip=["count():Q"],
        )
        .properties(title="HCC Payment Risk Score Distribution", width=600, height=300)
    )

    risk_scatter = (
        alt.Chart(risk.to_pandas())
        .mark_circle(opacity=0.6)
        .encode(
            x=alt.X("age:Q", title="Age"),
            y=alt.Y("payment_risk_score:Q", title="Payment Risk Score"),
            color=alt.Color("sex:N", scale=alt.Scale(range=PALETTE)),
            size=alt.Size(
                "member_months:Q", scale=alt.Scale(range=[20, 200]), legend=None
            ),
            tooltip=["person_id:N", "age:Q", "sex:N", "payment_risk_score:Q"],
        )
        .properties(title="Risk Score vs Age", width=600, height=300)
    )

    mo.vstack([risk_hist, risk_scatter])
    return (risk,)


@app.cell(hide_code=True)
def _(mo, risk):
    mo.md(f"""
    ### Risk Score Summary

    | Metric | Value |
    |--------|-------|
    | Patients scored | **{risk.height}** |
    | Mean risk score | **{risk["payment_risk_score"].mean():.3f}** |
    | Median risk score | **{risk["payment_risk_score"].median():.3f}** |
    | Std dev | **{risk["payment_risk_score"].std():.3f}** |
    | Min | **{risk["payment_risk_score"].min():.3f}** |
    | Max | **{risk["payment_risk_score"].max():.3f}** |
    """)
    return


@app.cell(hide_code=True)
def _(mo):
    mo.md("""
    ## Financial PMPM Trends
    """)
    return


@app.cell(hide_code=True)
def _(PALETTE, alt, mo, pl, q):
    pmpm = q("""
        SELECT year_month,
               member_months,
               total_paid / member_months AS total_pmpm,
               inpatient_paid / member_months AS inpatient_pmpm,
               outpatient_paid / member_months AS outpatient_pmpm,
               office_based_paid / member_months AS office_pmpm,
               pharmacy_paid / member_months AS pharmacy_pmpm,
               emergency_department_paid / member_months AS ed_pmpm
        FROM financial_pmpm.pmpm_payer
        WHERE member_months > 50
        ORDER BY year_month
    """)

    pmpm_long = pmpm.unpivot(
        index=["year_month", "member_months"],
        on=[
            "inpatient_pmpm",
            "outpatient_pmpm",
            "office_pmpm",
            "pharmacy_pmpm",
            "ed_pmpm",
        ],
        variable_name="category",
        value_name="pmpm",
    ).with_columns(
        pl.col("category")
        .str.replace("_pmpm", "")
        .str.replace("_", " ")
        .str.to_titlecase()
    )

    area_chart = (
        alt.Chart(pmpm_long.to_pandas())
        .mark_area(opacity=0.7)
        .encode(
            x=alt.X("year_month:T", title="Month"),
            y=alt.Y("pmpm:Q", stack=True, title="PMPM ($)"),
            color=alt.Color(
                "category:N", title="Category", scale=alt.Scale(range=PALETTE)
            ),
            tooltip=[
                "year_month:T",
                "category:N",
                alt.Tooltip("pmpm:Q", format="$.2f"),
            ],
        )
        .properties(title="PMPM by Service Category (Stacked)", width=800, height=350)
    )

    total_line = (
        alt.Chart(pmpm.to_pandas())
        .mark_line(color=PALETTE[1], strokeWidth=2, point=True)
        .encode(
            x=alt.X("year_month:T", title="Month"),
            y=alt.Y("total_pmpm:Q", title="Total PMPM ($)"),
            tooltip=["year_month:T", alt.Tooltip("total_pmpm:Q", format="$.2f")],
        )
        .properties(title="Total PMPM Trend", width=800, height=250)
    )

    mo.vstack([area_chart, total_line])
    return


@app.cell(hide_code=True)
def _(mo):
    mo.md("""
    ## Quality Measures Performance
    """)
    return


@app.cell(hide_code=True)
def _(PALETTE, alt, mo, pl, q):
    qm = q("""
        SELECT measure_id, measure_name, denominator_sum, numerator_sum,
               exclusion_sum, performance_rate
        FROM quality_measures.summary_counts
        ORDER BY measure_name
    """)

    # Convert Decimal columns to float for Altair compatibility
    qm = qm.with_columns(
        [
            pl.col("performance_rate").cast(pl.Float64),
            pl.col("denominator_sum").cast(pl.Float64),
            pl.col("numerator_sum").cast(pl.Float64),
            pl.col("exclusion_sum").cast(pl.Float64),
        ]
    )

    qm_chart = (
        alt.Chart(qm.to_pandas())
        .mark_bar()
        .encode(
            x=alt.X(
                "performance_rate:Q",
                title="Performance Rate (%)",
                scale=alt.Scale(domain=[0, 100]),
            ),
            y=alt.Y("measure_name:N", sort="-x", title=None),
            color=alt.condition(
                alt.datum.performance_rate > 50,
                alt.value(PALETTE[2]),
                alt.value(PALETTE[1]),
            ),
            tooltip=[
                "measure_id:N",
                "measure_name:N",
                alt.Tooltip("performance_rate:Q", format=".1f"),
                "denominator_sum:Q",
                "numerator_sum:Q",
            ],
        )
        .properties(title="Quality Measure Performance Rates", width=700, height=250)
    )

    mo.vstack(
        [
            qm_chart,
            mo.ui.table(qm, label="Quality Measures Detail"),
        ]
    )
    return


@app.cell(hide_code=True)
def _(mo):
    mo.md("""
    ## Emergency Department Classification
    """)
    return


@app.cell(hide_code=True)
def _(PALETTE, alt, mo, q):
    ed = q("""
        SELECT ed_classification_description AS ed_classification,
               count(*) AS visits,
               round(avg(charge_amount), 2) AS avg_charge,
               round(avg(paid_amount), 2) AS avg_paid
        FROM ed_classification.summary
        GROUP BY ed_classification_description
        ORDER BY visits DESC
    """)

    ed_pie = (
        alt.Chart(ed.to_pandas())
        .mark_arc(innerRadius=60)
        .encode(
            theta=alt.Theta("visits:Q"),
            color=alt.Color(
                "ed_classification:N",
                title="Classification",
                scale=alt.Scale(range=PALETTE),
            ),
            tooltip=[
                "ed_classification:N",
                "visits:Q",
                alt.Tooltip("avg_charge:Q", format="$,.0f"),
                alt.Tooltip("avg_paid:Q", format="$,.0f"),
            ],
        )
        .properties(title="ED Visits by Classification", width=400, height=350)
    )

    ed_cost = (
        alt.Chart(ed.to_pandas())
        .mark_bar()
        .encode(
            x=alt.X(
                "ed_classification:N",
                sort="-y",
                title=None,
                axis=alt.Axis(labelAngle=-30),
            ),
            y=alt.Y("avg_paid:Q", title="Avg Paid ($)"),
            color=alt.Color(
                "ed_classification:N", scale=alt.Scale(range=PALETTE), legend=None
            ),
            tooltip=["ed_classification:N", alt.Tooltip("avg_paid:Q", format="$,.0f")],
        )
        .properties(
            title="Average ED Paid Amount by Classification", width=400, height=350
        )
    )

    mo.hstack([ed_pie, ed_cost], justify="start", gap=2)
    return


@app.cell(hide_code=True)
def _(PALETTE, alt, mo, q):
    ed_dx = q("""
        SELECT primary_diagnosis_description AS diagnosis, count(*) AS visits
        FROM ed_classification.summary
        GROUP BY diagnosis
        ORDER BY visits DESC
        LIMIT 15
    """)

    ed_dx_chart = (
        alt.Chart(ed_dx.to_pandas())
        .mark_bar()
        .encode(
            x=alt.X("visits:Q", title="ED Visits"),
            y=alt.Y("diagnosis:N", sort="-x", title=None),
            color=alt.value(PALETTE[0]),
            tooltip=["diagnosis:N", "visits:Q"],
        )
        .properties(title="Top 15 ED Diagnoses", width=700, height=400)
    )

    ed_dx_chart
    return


@app.cell(hide_code=True)
def _(mo):
    mo.md("""
    ## Readmissions Analysis
    """)
    return


@app.cell(hide_code=True)
def _(mo, q):
    readmit = q("""
        SELECT count(*) AS index_admissions,
               sum(CASE WHEN readmit_30_flag = 1 THEN 1 ELSE 0 END) AS readmissions,
               round(100.0 * sum(CASE WHEN readmit_30_flag = 1 THEN 1 ELSE 0 END) / count(*), 1) AS readmit_rate_pct,
               round(avg(length_of_stay), 1) AS avg_index_los,
               round(avg(CASE WHEN readmit_30_flag = 1 THEN days_to_readmit END), 1) AS avg_days_to_readmit
        FROM readmissions.readmission_summary
    """)

    r = readmit.row(0, named=True)
    mo.md(f"""
    ### Readmission Summary

    | Metric | Value |
    |--------|-------|
    | Index admissions | **{r["index_admissions"]}** |
    | Readmissions | **{r["readmissions"]}** |
    | Readmission rate | **{r["readmit_rate_pct"]}%** |
    | Avg index LOS | **{r["avg_index_los"]} days** |
    | Avg days to readmit | **{r["avg_days_to_readmit"]} days** |
    """)
    return


@app.cell(hide_code=True)
def _(PALETTE, alt, mo, q):
    readmit_cohort = q("""
        SELECT specialty_cohort AS cohort,
               count(*) AS admissions,
               sum(CASE WHEN readmit_30_flag = 1 THEN 1 ELSE 0 END) AS readmissions,
               round(100.0 * sum(CASE WHEN readmit_30_flag = 1 THEN 1 ELSE 0 END) / count(*), 1) AS rate
        FROM readmissions.readmission_summary
        GROUP BY cohort
        HAVING count(*) >= 5
        ORDER BY rate DESC
    """)

    cohort_chart = (
        alt.Chart(readmit_cohort.to_pandas())
        .mark_bar()
        .encode(
            x=alt.X("rate:Q", title="Readmission Rate (%)"),
            y=alt.Y("cohort:N", sort="-x", title=None),
            color=alt.condition(
                alt.datum.rate > 20,
                alt.value(PALETTE[1]),
                alt.value(PALETTE[2]),
            ),
            tooltip=[
                "cohort:N",
                "admissions:Q",
                "readmissions:Q",
                alt.Tooltip("rate:Q", format=".1f"),
            ],
        )
        .properties(title="Readmission Rate by Specialty Cohort", width=600, height=250)
    )

    cohort_chart
    return


@app.cell(hide_code=True)
def _(mo):
    mo.md("""
    ## Pharmacy: Brand vs Generic Opportunity
    """)
    return


@app.cell(hide_code=True)
def _(PALETTE, alt, mo, pl, q):
    brand_gen = q("""
        SELECT brand_vs_generic, count(*) AS claims,
               round(sum(paid_amount), 2) AS total_paid
        FROM pharmacy.pharmacy_claim_expanded
        WHERE brand_vs_generic IS NOT NULL
        GROUP BY brand_vs_generic
    """)

    # Convert Decimal columns to float for Altair compatibility
    brand_gen = brand_gen.with_columns(
        [
            pl.col("total_paid").cast(pl.Float64),
            pl.col("claims").cast(pl.Float64),
        ]
    )

    bg_chart = (
        alt.Chart(brand_gen.to_pandas())
        .mark_bar()
        .encode(
            x=alt.X("brand_vs_generic:N", title=None),
            y=alt.Y("total_paid:Q", title="Total Paid ($)"),
            color=alt.Color(
                "brand_vs_generic:N",
                scale=alt.Scale(
                    domain=["brand", "generic"], range=[PALETTE[7], PALETTE[2]]
                ),
                legend=None,
            ),
            tooltip=[
                "brand_vs_generic:N",
                "claims:Q",
                alt.Tooltip("total_paid:Q", format="$,.0f"),
            ],
        )
        .properties(width=300, height=300)
    )

    top_drugs = q("""
        SELECT ndc_description, brand_vs_generic,
               count(*) AS claims,
               round(sum(paid_amount), 2) AS total_paid
        FROM pharmacy.pharmacy_claim_expanded
        WHERE ndc_description IS NOT NULL
        GROUP BY ndc_description, brand_vs_generic
        ORDER BY total_paid DESC
        LIMIT 20
    """)

    # Convert Decimal columns to float for Altair compatibility
    top_drugs = top_drugs.with_columns(
        [
            pl.col("total_paid").cast(pl.Float64),
            pl.col("claims").cast(pl.Float64),
        ]
    )

    mo.hstack(
        [
            bg_chart,
            mo.ui.table(top_drugs, label="Top Drugs by Spend"),
        ],
        justify="start",
        gap=2,
    )
    return


@app.cell(hide_code=True)
def _(mo):
    mo.md("""
    ## AHRQ Prevention Quality Indicators
    """)
    return


@app.cell(hide_code=True)
def _(PALETTE, alt, mo, q):
    pqi = q("""
        SELECT r.pqi_number, m.pqi_name AS pqi_description,
               r.rate_per_100_thousand, r.num_count, r.denom_count
        FROM ahrq_measures.pqi_rate r
        JOIN ahrq_measures._value_set_pqi_measures m ON r.pqi_number = m.pqi_number
        WHERE r.rate_per_100_thousand > 0
        ORDER BY r.rate_per_100_thousand DESC
    """)

    pqi_chart = (
        alt.Chart(pqi.to_pandas())
        .mark_bar()
        .encode(
            x=alt.X("rate_per_100_thousand:Q", title="Rate per 100,000"),
            y=alt.Y("pqi_description:N", sort="-x", title=None),
            color=alt.value(PALETTE[0]),
            tooltip=[
                "pqi_number:N",
                "pqi_description:N",
                alt.Tooltip("rate_per_100_thousand:Q", format=",.0f"),
                "num_count:Q",
                "denom_count:Q",
            ],
        )
        .properties(title="PQI Observed Rates", width=700, height=300)
    )

    pqi_chart
    return


if __name__ == "__main__":
    app.run()