{"spec_id":"survival-kaplan-meier","library":"altair","language":"python","code":"\"\"\" anyplot.ai\nsurvival-kaplan-meier: Kaplan-Meier Survival Plot\nLibrary: altair 6.1.0 | Python 3.13.13\nQuality: 90/100 | Updated: 2026-05-11\n\"\"\"\n\nimport os\n\nimport altair as alt\nimport numpy as np\nimport pandas as pd\n\n\n# Theme tokens\nTHEME = os.getenv(\"ANYPLOT_THEME\", \"light\")\nPAGE_BG = \"#FAF8F1\" if THEME == \"light\" else \"#1A1A17\"\nELEVATED_BG = \"#FFFDF6\" if THEME == \"light\" else \"#242420\"\nINK = \"#1A1A17\" if THEME == \"light\" else \"#F0EFE8\"\nINK_SOFT = \"#4A4A44\" if THEME == \"light\" else \"#B8B7B0\"\n\n# Okabe-Ito palette\nIMPRINT = [\"#009E73\", \"#C475FD\", \"#4467A3\", \"#BD8233\", \"#AE3030\", \"#2ABCCD\", \"#954477\"]\n\n# Data - Clinical trial with two treatment groups\nnp.random.seed(42)\n\nn_per_group = 80\n\n# Treatment A (better survival)\ntime_a = np.random.exponential(scale=24, size=n_per_group)\ntime_a = np.clip(time_a, 1, 36)\nevent_a = np.random.binomial(1, 0.65, size=n_per_group)\n\n# Treatment B (standard)\ntime_b = np.random.exponential(scale=16, size=n_per_group)\ntime_b = np.clip(time_b, 1, 36)\nevent_b = np.random.binomial(1, 0.75, size=n_per_group)\n\n# Combine into dataframe\ndf = pd.DataFrame(\n    {\n        \"time\": np.concatenate([time_a, time_b]),\n        \"event\": np.concatenate([event_a, event_b]),\n        \"group\": [\"Treatment A\"] * n_per_group + [\"Treatment B\"] * n_per_group,\n    }\n)\n\n# Kaplan-Meier estimation (inline - no functions)\nkm_data = []\n\nfor group_name in [\"Treatment A\", \"Treatment B\"]:\n    mask = df[\"group\"] == group_name\n    time = df.loc[mask, \"time\"].values\n    event = df.loc[mask, \"event\"].values\n\n    # Sort by time\n    order = np.argsort(time)\n    time = time[order]\n    event = event[order]\n\n    # Get unique event times\n    unique_times = np.unique(time[event == 1])\n\n    # Calculate survival at each time point\n    survival = 1.0\n    times = [0]\n    survivals = [1.0]\n    ci_lower = [1.0]\n    ci_upper = [1.0]\n    var_sum = 0\n\n    for t in unique_times:\n        at_risk = np.sum(time >= t)\n        events = np.sum((time == t) & (event == 1))\n\n        if at_risk > 0:\n            survival *= (at_risk - events) / at_risk\n            if at_risk > events:\n                var_sum += events / (at_risk * (at_risk - events))\n\n        times.append(t)\n        survivals.append(survival)\n\n        se = survival * np.sqrt(var_sum) if var_sum > 0 else 0\n        ci_lower.append(max(0, survival - 1.96 * se))\n        ci_upper.append(min(1, survival + 1.96 * se))\n\n    # Extend to max time\n    max_time = time.max()\n    times.append(max_time)\n    survivals.append(survival)\n    ci_lower.append(ci_lower[-1])\n    ci_upper.append(ci_upper[-1])\n\n    for i in range(len(times)):\n        km_data.append(\n            {\n                \"Time (Months)\": times[i],\n                \"Survival Probability\": survivals[i],\n                \"CI Lower\": ci_lower[i],\n                \"CI Upper\": ci_upper[i],\n                \"Group\": group_name,\n            }\n        )\n\nkm_df = pd.DataFrame(km_data)\n\n# Get censored observations for tick marks\ncensored = df[df[\"event\"] == 0].copy()\ncensored_marks = []\nfor _, row in censored.iterrows():\n    mask = (km_df[\"Group\"] == row[\"group\"]) & (km_df[\"Time (Months)\"] <= row[\"time\"])\n    if mask.any():\n        surv_at_censor = km_df.loc[mask, \"Survival Probability\"].iloc[-1]\n        censored_marks.append(\n            {\"Time (Months)\": row[\"time\"], \"Survival Probability\": surv_at_censor, \"Group\": row[\"group\"]}\n        )\n\ncensored_df = pd.DataFrame(censored_marks) if censored_marks else pd.DataFrame()\n\n# Define colors using Okabe-Ito palette\ncolor_scale = alt.Scale(domain=[\"Treatment A\", \"Treatment B\"], range=[IMPRINT[0], IMPRINT[1]])\n\n# Step line for survival curves\nsurvival_line = (\n    alt.Chart(km_df)\n    .mark_line(interpolate=\"step-after\", strokeWidth=4)\n    .encode(\n        x=alt.X(\"Time (Months):Q\", scale=alt.Scale(domain=[0, 38]), title=\"Time (Months)\"),\n        y=alt.Y(\"Survival Probability:Q\", scale=alt.Scale(domain=[0, 1.0]), title=\"Survival Probability\"),\n        color=alt.Color(\"Group:N\", scale=color_scale),\n    )\n)\n\n# Confidence interval bands\nci_band = (\n    alt.Chart(km_df)\n    .mark_area(interpolate=\"step-after\", opacity=0.25)\n    .encode(\n        x=alt.X(\"Time (Months):Q\"),\n        y=alt.Y(\"CI Lower:Q\", title=\"\"),\n        y2=alt.Y2(\"CI Upper:Q\"),\n        color=alt.Color(\"Group:N\", scale=color_scale, legend=None),\n    )\n)\n\n# Censored observation marks\nif not censored_df.empty:\n    censor_marks = (\n        alt.Chart(censored_df)\n        .mark_tick(thickness=3, size=25)\n        .encode(\n            x=alt.X(\"Time (Months):Q\"),\n            y=alt.Y(\"Survival Probability:Q\", title=\"\"),\n            color=alt.Color(\"Group:N\", scale=color_scale, legend=None),\n        )\n    )\n    chart_layers = ci_band + survival_line + censor_marks\nelse:\n    chart_layers = ci_band + survival_line\n\n# Create final chart with theme-adaptive styling\nchart = (\n    chart_layers.properties(\n        width=1600,\n        height=900,\n        background=PAGE_BG,\n        title=alt.Title(\"survival-kaplan-meier · altair · anyplot.ai\", fontSize=28, anchor=\"middle\", offset=20),\n    )\n    .configure_view(fill=PAGE_BG, strokeWidth=0)\n    .configure_axis(\n        domainColor=INK_SOFT,\n        tickColor=INK_SOFT,\n        gridColor=INK,\n        gridOpacity=0.10,\n        labelFontSize=18,\n        titleFontSize=22,\n        labelColor=INK_SOFT,\n        titleColor=INK,\n    )\n    .configure_title(color=INK)\n    .configure_legend(\n        fillColor=ELEVATED_BG,\n        strokeColor=INK_SOFT,\n        labelColor=INK_SOFT,\n        titleColor=INK,\n        titleFontSize=20,\n        labelFontSize=18,\n        strokeWidth=1,\n    )\n)\n\n# Save outputs\nchart.save(f\"plot-{THEME}.png\", scale_factor=3.0)\nchart.save(f\"plot-{THEME}.html\")\n"}