Oncology swimmer plot

Python Code

"""Oncology swimmer plot — Rotterdam breast cancer dataset, 50 patient sample."""


import numpy as np
import pandas as pd
import plotly.graph_objects as go


ROTTERDAM_URL = (
    "https://vincentarelbundock.github.io/Rdatasets/csv/OncoDataSets/RotterdamBreastCancer_df.csv"
)

N_PATIENTS = 50


def generate():
    print("fetching Rotterdam breast cancer dataset …")
    df = None
    try:
        df = fetch_csv(ROTTERDAM_URL)
        df.columns = [c.strip() for c in df.columns]
        print(f"  cols: {list(df.columns)}, shape: {df.shape}")
    except Exception as e:
        print(f"  fetch failed ({e}), using synthetic swimmer data")

    if df is not None:
        # Identify required columns
        id_col    = next((c for c in df.columns if "pid" in c.lower() or "id" in c.lower()), None)
        rtime_col = next((c for c in df.columns if "rtime" in c.lower()), None)
        recur_col = next((c for c in df.columns if "recur" in c.lower()), None)
        dtime_col = next((c for c in df.columns if "dtime" in c.lower()), None)
        death_col = next((c for c in df.columns if "death" in c.lower()), None)
        chemo_col = next((c for c in df.columns if "chemo" in c.lower()), None)

        print(f"  id={id_col}, rtime={rtime_col}, recur={recur_col}, "
              f"dtime={dtime_col}, death={death_col}, chemo={chemo_col}")

        if all(c is not None for c in [rtime_col, recur_col, dtime_col, chemo_col]):
            for col in [rtime_col, recur_col, dtime_col, chemo_col]:
                df[col] = pd.to_numeric(df[col], errors="coerce")
            df = df.dropna(subset=[rtime_col, recur_col, dtime_col, chemo_col])

            # Subsample 50 patients, balanced by chemo
            chemo1 = df[df[chemo_col] == 1].sample(n=min(25, (df[chemo_col] == 1).sum()),
                                                     random_state=43)
            chemo0 = df[df[chemo_col] == 0].sample(n=min(25, (df[chemo_col] == 0).sum()),
                                                     random_state=43)
            df = pd.concat([chemo1, chemo0]).reset_index(drop=True)
            df = df.sort_values(dtime_col, ascending=False).reset_index(drop=True)
            df["patient_idx"] = np.arange(len(df))
            print(f"  sampled {len(df)} patients")
        else:
            print(f"  missing required columns")
            df = None

    if df is None:
        # Synthetic Rotterdam-like data
        rng = np.random.default_rng(43)
        n = N_PATIENTS
        chemo = rng.integers(0, 2, n)
        # Chemo patients live slightly longer on average
        dtime = rng.exponential(scale=np.where(chemo == 1, 2200, 1800), size=n).astype(int)
        dtime = np.clip(dtime, 30, 5000)
        death = (rng.random(n) < 0.55).astype(int)
        recur = (rng.random(n) < 0.40).astype(int)
        rtime = np.where(recur, rng.uniform(0.2, 0.9, n) * dtime, dtime).astype(int)

        df = pd.DataFrame({
            "patient_idx": np.arange(n),
            "dtime": dtime,
            "death": death,
            "rtime": rtime,
            "recur": recur,
            "chemo": chemo,
        })
        df = df.sort_values("dtime", ascending=False).reset_index(drop=True)
        df["patient_idx"] = np.arange(n)

        dtime_col = "dtime"
        rtime_col = "rtime"
        recur_col = "recur"
        death_col = "death"
        chemo_col = "chemo"

    # Ensure columns are set if we used real data
    if dtime_col not in ["dtime"]:
        df = df.rename(columns={
            dtime_col: "dtime", rtime_col: "rtime",
            recur_col: "recur", chemo_col: "chemo",
        })
        if death_col and death_col != "death":
            df = df.rename(columns={death_col: "death"})
        dtime_col = "dtime"; rtime_col = "rtime"
        recur_col = "recur"; chemo_col = "chemo"

    fig = go.Figure()

    # Bars grouped by chemo status
    for chemo_val, color, label in [
        (1, VIOLET, "Chemotherapy"),
        (0, TEAL,   "No chemotherapy"),
    ]:
        sub = df[df[chemo_col] == chemo_val]
        if len(sub) == 0:
            continue
        fig.add_trace(go.Bar(
            x=sub[dtime_col].values,
            y=sub["patient_idx"].values,
            base=0,
            orientation="h",
            name=label,
            marker=dict(
                color=color,
                opacity=0.6,
                line=dict(color=color, width=0.5),
            ),
            hovertemplate=(
                f"Patient %{{y}} ({label})<br>"
                "Follow-up: %{x} days<extra></extra>"
            ),
            width=0.7,
        ))

    # Death markers (triangle-right) at dtime
    if "death" in df.columns:
        died = df[df["death"] == 1]
        if len(died):
            fig.add_trace(go.Scatter(
                x=died[dtime_col].values,
                y=died["patient_idx"].values,
                mode="markers",
                name="Death",
                marker=dict(color="#333333", size=7, symbol="triangle-right"),
                hovertemplate="Patient %{y}: died at %{x} days<extra></extra>",
            ))

    # Recurrence markers (diamond) at rtime where recur==1
    recurred = df[df[recur_col] == 1]
    if len(recurred):
        fig.add_trace(go.Scatter(
            x=recurred[rtime_col].values,
            y=recurred["patient_idx"].values,
            mode="markers",
            name="Recurrence",
            marker=dict(color=PINK, size=9, symbol="diamond",
                        line=dict(color="white", width=0.5)),
            hovertemplate="Patient %{y}: recurrence at %{x} days<extra></extra>",
        ))

    fig.update_layout(
        xaxis=dict(title="Follow-up (days)", showgrid=False),
        yaxis=dict(
            title="Patient",
            tickvals=df["patient_idx"].values[::5],
            ticktext=[str(i) for i in df["patient_idx"].values[::5]],
            showgrid=False,
        ),
        barmode="overlay",
        legend=dict(orientation="h", y=1.08),
        margin=dict(t=50, b=50, l=60, r=40),
        height=600,
    )
    return fig


if __name__ == "__main__":
    generate()

Made with Plotly