{"spec_id":"survival-kaplan-meier","library":"plotly","language":"python","code":"\"\"\" anyplot.ai\nsurvival-kaplan-meier: Kaplan-Meier Survival Plot\nLibrary: plotly 6.7.0 | Python 3.13.13\nQuality: 92/100 | Updated: 2026-05-11\n\"\"\"\n\nimport os\nimport sys\n\n\n# Fix import issue: remove script directory from path before importing plotly\ntry:\n    _script_dir = os.path.dirname(os.path.abspath(__file__))\nexcept NameError:\n    _script_dir = os.getcwd()\n\nif _script_dir in sys.path:\n    sys.path.remove(_script_dir)\nif \"\" in sys.path:\n    sys.path.remove(\"\")\n\nimport numpy as np\nimport plotly.graph_objects as go\n\n\n# Ensure we save to the script's directory\nSCRIPT_DIR = _script_dir if \"_script_dir\" in locals() else os.path.dirname(os.path.abspath(__file__))\n\n\ndef hex_to_rgba(hex_color, alpha):\n    hex_color = hex_color.lstrip(\"#\")\n    r = int(hex_color[0:2], 16)\n    g = int(hex_color[2:4], 16)\n    b = int(hex_color[4:6], 16)\n    return f\"rgba({r}, {g}, {b}, {alpha})\"\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\"\nGRID = \"rgba(26,26,23,0.10)\" if THEME == \"light\" else \"rgba(240,239,232,0.10)\"\n\n# Okabe-Ito palette\nCOLOR_A = \"#009E73\"  # First series (bluish green)\nCOLOR_B = \"#C475FD\"  # Second series (vermillion)\n\n# Data - Clinical trial with two treatment groups\nnp.random.seed(42)\n\n# Treatment A (experimental drug) - better survival\nn_a = 80\ntime_a = np.random.exponential(scale=24, size=n_a)\ntime_a = np.clip(time_a, 1, 48)\nevent_a = np.random.binomial(1, 0.7, size=n_a)\ncensored_idx = time_a > 36\nevent_a[censored_idx] = 0\n\n# Treatment B (standard care) - worse survival\nn_b = 80\ntime_b = np.random.exponential(scale=16, size=n_b)\ntime_b = np.clip(time_b, 1, 48)\nevent_b = np.random.binomial(1, 0.8, size=n_b)\ncensored_idx = time_b > 36\nevent_b[censored_idx] = 0\n\n# Kaplan-Meier estimator for Treatment A\norder_a = np.argsort(time_a)\ntime_a_sorted = time_a[order_a]\nevent_a_sorted = event_a[order_a]\nunique_times_a = np.unique(time_a_sorted[event_a_sorted == 1])\n\nkm_times_a = []\nkm_survival_a = []\nkm_ci_lower_a = []\nkm_ci_upper_a = []\nS_a = 1.0\nvar_sum_a = 0.0\n\nfor t in unique_times_a:\n    at_risk = np.sum(time_a_sorted >= t)\n    events = np.sum((time_a_sorted == t) & (event_a_sorted == 1))\n    if at_risk > 0:\n        S_a = S_a * (1 - events / at_risk)\n        if at_risk > events:\n            var_sum_a += events / (at_risk * (at_risk - events))\n        if S_a > 0 and S_a < 1 and var_sum_a > 0:\n            se_log = np.sqrt(var_sum_a) / abs(np.log(S_a))\n            lower = S_a ** np.exp(1.96 * se_log)\n            upper = S_a ** np.exp(-1.96 * se_log)\n        else:\n            lower = max(0, S_a - 0.1)\n            upper = min(1, S_a + 0.1)\n        km_times_a.append(t)\n        km_survival_a.append(S_a)\n        km_ci_lower_a.append(max(0, lower))\n        km_ci_upper_a.append(min(1, upper))\n\nkm_times_a = np.array(km_times_a)\nkm_survival_a = np.array(km_survival_a)\nkm_ci_lower_a = np.array(km_ci_lower_a)\nkm_ci_upper_a = np.array(km_ci_upper_a)\n\n# Kaplan-Meier estimator for Treatment B\norder_b = np.argsort(time_b)\ntime_b_sorted = time_b[order_b]\nevent_b_sorted = event_b[order_b]\nunique_times_b = np.unique(time_b_sorted[event_b_sorted == 1])\n\nkm_times_b = []\nkm_survival_b = []\nkm_ci_lower_b = []\nkm_ci_upper_b = []\nS_b = 1.0\nvar_sum_b = 0.0\n\nfor t in unique_times_b:\n    at_risk = np.sum(time_b_sorted >= t)\n    events = np.sum((time_b_sorted == t) & (event_b_sorted == 1))\n    if at_risk > 0:\n        S_b = S_b * (1 - events / at_risk)\n        if at_risk > events:\n            var_sum_b += events / (at_risk * (at_risk - events))\n        if S_b > 0 and S_b < 1 and var_sum_b > 0:\n            se_log = np.sqrt(var_sum_b) / abs(np.log(S_b))\n            lower = S_b ** np.exp(1.96 * se_log)\n            upper = S_b ** np.exp(-1.96 * se_log)\n        else:\n            lower = max(0, S_b - 0.1)\n            upper = min(1, S_b + 0.1)\n        km_times_b.append(t)\n        km_survival_b.append(S_b)\n        km_ci_lower_b.append(max(0, lower))\n        km_ci_upper_b.append(min(1, upper))\n\nkm_times_b = np.array(km_times_b)\nkm_survival_b = np.array(km_survival_b)\nkm_ci_lower_b = np.array(km_ci_lower_b)\nkm_ci_upper_b = np.array(km_ci_upper_b)\n\n# Create step coordinates for Treatment A\nx_step_a = [0.0]\ny_step_a = [1.0]\ny_lower_a = [1.0]\ny_upper_a = [1.0]\nfor i, t in enumerate(km_times_a):\n    prev_surv = 1.0 if i == 0 else km_survival_a[i - 1]\n    prev_lower = 1.0 if i == 0 else km_ci_lower_a[i - 1]\n    prev_upper = 1.0 if i == 0 else km_ci_upper_a[i - 1]\n    x_step_a.extend([t, t])\n    y_step_a.extend([prev_surv, km_survival_a[i]])\n    y_lower_a.extend([prev_lower, km_ci_lower_a[i]])\n    y_upper_a.extend([prev_upper, km_ci_upper_a[i]])\nx_step_a = np.array(x_step_a)\ny_step_a = np.array(y_step_a)\ny_lower_a = np.array(y_lower_a)\ny_upper_a = np.array(y_upper_a)\n\n# Create step coordinates for Treatment B\nx_step_b = [0.0]\ny_step_b = [1.0]\ny_lower_b = [1.0]\ny_upper_b = [1.0]\nfor i, t in enumerate(km_times_b):\n    prev_surv = 1.0 if i == 0 else km_survival_b[i - 1]\n    prev_lower = 1.0 if i == 0 else km_ci_lower_b[i - 1]\n    prev_upper = 1.0 if i == 0 else km_ci_upper_b[i - 1]\n    x_step_b.extend([t, t])\n    y_step_b.extend([prev_surv, km_survival_b[i]])\n    y_lower_b.extend([prev_lower, km_ci_lower_b[i]])\n    y_upper_b.extend([prev_upper, km_ci_upper_b[i]])\nx_step_b = np.array(x_step_b)\ny_step_b = np.array(y_step_b)\ny_lower_b = np.array(y_lower_b)\ny_upper_b = np.array(y_upper_b)\n\n# Create figure\nfig = go.Figure()\n\n# Treatment A - Confidence interval band\nfig.add_trace(\n    go.Scatter(\n        x=np.concatenate([x_step_a, x_step_a[::-1]]),\n        y=np.concatenate([y_upper_a, y_lower_a[::-1]]),\n        fill=\"toself\",\n        fillcolor=hex_to_rgba(COLOR_A, 0.15),\n        line=dict(color=\"rgba(0,0,0,0)\"),\n        showlegend=False,\n        hoverinfo=\"skip\",\n        name=\"CI A\",\n    )\n)\n\n# Treatment B - Confidence interval band\nfig.add_trace(\n    go.Scatter(\n        x=np.concatenate([x_step_b, x_step_b[::-1]]),\n        y=np.concatenate([y_upper_b, y_lower_b[::-1]]),\n        fill=\"toself\",\n        fillcolor=hex_to_rgba(COLOR_B, 0.15),\n        line=dict(color=\"rgba(0,0,0,0)\"),\n        showlegend=False,\n        hoverinfo=\"skip\",\n        name=\"CI B\",\n    )\n)\n\n# Treatment A - Survival curve\nfig.add_trace(\n    go.Scatter(\n        x=x_step_a,\n        y=y_step_a,\n        mode=\"lines\",\n        line=dict(color=COLOR_A, width=4),\n        name=\"Treatment A (Experimental)\",\n        hovertemplate=\"Time: %{x:.1f} months<br>Survival: %{y:.1%}<extra></extra>\",\n    )\n)\n\n# Treatment B - Survival curve\nfig.add_trace(\n    go.Scatter(\n        x=x_step_b,\n        y=y_step_b,\n        mode=\"lines\",\n        line=dict(color=COLOR_B, width=4),\n        name=\"Treatment B (Standard Care)\",\n        hovertemplate=\"Time: %{x:.1f} months<br>Survival: %{y:.1%}<extra></extra>\",\n    )\n)\n\n# Censored observations - Treatment A\ncensored_times_a = time_a[event_a == 0]\ncensored_surv_a = []\nfor t in censored_times_a:\n    idx = np.searchsorted(km_times_a, t)\n    if idx > 0:\n        censored_surv_a.append(km_survival_a[idx - 1])\n    else:\n        censored_surv_a.append(1.0)\n\nfig.add_trace(\n    go.Scatter(\n        x=censored_times_a,\n        y=censored_surv_a,\n        mode=\"markers\",\n        marker=dict(symbol=\"line-ns\", size=14, line=dict(width=3, color=COLOR_A)),\n        name=\"Censored\",\n        legendgroup=\"censored\",\n        hovertemplate=\"Censored at: %{x:.1f} months<extra></extra>\",\n    )\n)\n\n# Censored observations - Treatment B\ncensored_times_b = time_b[event_b == 0]\ncensored_surv_b = []\nfor t in censored_times_b:\n    idx = np.searchsorted(km_times_b, t)\n    if idx > 0:\n        censored_surv_b.append(km_survival_b[idx - 1])\n    else:\n        censored_surv_b.append(1.0)\n\nfig.add_trace(\n    go.Scatter(\n        x=censored_times_b,\n        y=censored_surv_b,\n        mode=\"markers\",\n        marker=dict(symbol=\"line-ns\", size=14, line=dict(width=3, color=COLOR_B)),\n        name=\"Censored\",\n        legendgroup=\"censored\",\n        showlegend=False,\n        hovertemplate=\"Censored at: %{x:.1f} months<extra></extra>\",\n    )\n)\n\n# Calculate median survival (time when S = 0.5)\nmedian_a = np.nan\nfor i, s in enumerate(km_survival_a):\n    if s <= 0.5:\n        median_a = km_times_a[i]\n        break\n\nmedian_b = np.nan\nfor i, s in enumerate(km_survival_b):\n    if s <= 0.5:\n        median_b = km_times_b[i]\n        break\n\n# Layout\nfig.update_layout(\n    title=dict(\n        text=\"survival-kaplan-meier · plotly · anyplot.ai\", font=dict(size=28, color=INK), x=0.5, xanchor=\"center\"\n    ),\n    xaxis=dict(\n        title=dict(text=\"Time (months)\", font=dict(size=22, color=INK)),\n        tickfont=dict(size=18, color=INK_SOFT),\n        range=[0, 50],\n        gridcolor=GRID,\n        showline=True,\n        linewidth=2,\n        linecolor=INK_SOFT,\n        zeroline=False,\n    ),\n    yaxis=dict(\n        title=dict(text=\"Survival Probability\", font=dict(size=22, color=INK)),\n        tickfont=dict(size=18, color=INK_SOFT),\n        tickformat=\".0%\",\n        range=[0, 1.05],\n        gridcolor=GRID,\n        showline=True,\n        linewidth=2,\n        linecolor=INK_SOFT,\n        zeroline=False,\n    ),\n    legend=dict(\n        x=0.98,\n        y=0.98,\n        xanchor=\"right\",\n        yanchor=\"top\",\n        font=dict(size=18, color=INK_SOFT),\n        bgcolor=ELEVATED_BG,\n        bordercolor=INK_SOFT,\n        borderwidth=1,\n    ),\n    paper_bgcolor=PAGE_BG,\n    plot_bgcolor=PAGE_BG,\n    margin=dict(l=100, r=80, t=100, b=100),\n)\n\n# Add median survival annotation\nmedian_a_str = f\"{median_a:.1f}\" if not np.isnan(median_a) else \"N/A\"\nmedian_b_str = f\"{median_b:.1f}\" if not np.isnan(median_b) else \"N/A\"\nmedian_text = f\"Median Survival:<br>Treatment A: {median_a_str} months<br>Treatment B: {median_b_str} months\"\n\nfig.add_annotation(\n    x=0.02,\n    y=0.02,\n    xref=\"paper\",\n    yref=\"paper\",\n    text=median_text,\n    showarrow=False,\n    font=dict(size=16, color=INK),\n    align=\"left\",\n    bgcolor=ELEVATED_BG,\n    bordercolor=INK_SOFT,\n    borderwidth=1,\n    borderpad=8,\n)\n\n# Save\nfig.write_image(os.path.join(SCRIPT_DIR, f\"plot-{THEME}.png\"), width=1600, height=900, scale=3)\nfig.write_html(os.path.join(SCRIPT_DIR, f\"plot-{THEME}.html\"), include_plotlyjs=\"cdn\")\n"}