{"spec_id":"survival-kaplan-meier","library":"letsplot","language":"python","code":"\"\"\" anyplot.ai\nsurvival-kaplan-meier: Kaplan-Meier Survival Plot\nLibrary: letsplot 4.9.0 | Python 3.13.13\nQuality: 94/100 | Updated: 2026-05-11\n\"\"\"\n\nimport os\nimport shutil\n\nimport numpy as np\nimport pandas as pd\nfrom lets_plot import (\n    LetsPlot,\n    aes,\n    element_line,\n    element_rect,\n    element_text,\n    geom_point,\n    geom_ribbon,\n    geom_step,\n    ggplot,\n    ggsave,\n    ggsize,\n    labs,\n    scale_color_manual,\n    scale_fill_manual,\n    theme,\n    theme_minimal,\n)\n\n\nLetsPlot.setup_html()\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\"\nINK_MUTED = \"#6B6A63\" if THEME == \"light\" else \"#A8A79F\"\n\n# Okabe-Ito palette (brand color #009E73 is first series)\nIMPRINT = [\"#009E73\", \"#C475FD\"]\n\n# Data - Simulated clinical trial survival data for two treatment groups\nnp.random.seed(42)\n\nn_per_group = 80\n\n\ndef generate_survival_data(n, hazard_rate, group_name):\n    \"\"\"Generate survival data with exponential distribution.\"\"\"\n    times = np.random.exponential(scale=1 / hazard_rate, size=n)\n    times = np.clip(times, 0, 36)  # Max follow-up 36 months\n    censored = times >= 36\n    times[censored] = 36\n    event = (~censored).astype(int)\n    # Add random censoring (20% of non-terminal events)\n    random_censor = np.random.random(n) < 0.2\n    event[random_censor] = 0\n    return pd.DataFrame({\"time\": times, \"event\": event, \"group\": group_name})\n\n\n# Treatment group (lower hazard = better survival)\ntreatment = generate_survival_data(n_per_group, hazard_rate=0.04, group_name=\"Treatment\")\n# Control group (higher hazard = worse survival)\ncontrol = generate_survival_data(n_per_group, hazard_rate=0.08, group_name=\"Control\")\ndf = pd.concat([treatment, control], ignore_index=True)\n\n\n# Kaplan-Meier estimator function\ndef kaplan_meier(time, event):\n    \"\"\"Compute Kaplan-Meier survival curve with confidence intervals.\"\"\"\n    df_km = pd.DataFrame({\"time\": time, \"event\": event}).sort_values(\"time\")\n    unique_times = np.sort(df_km[\"time\"].unique())\n    n_at_risk = len(df_km)\n    survival = 1.0\n    results = [{\"time\": 0, \"survival\": 1.0, \"ci_lower\": 1.0, \"ci_upper\": 1.0, \"n_at_risk\": n_at_risk}]\n    var_sum = 0\n\n    for t in unique_times:\n        at_time = df_km[df_km[\"time\"] == t]\n        d = at_time[\"event\"].sum()  # Number of events\n        n = n_at_risk  # Number at risk\n        if n > 0 and d > 0:\n            survival *= 1 - d / n\n            var_sum += d / (n * (n - d)) if n > d else 0\n        # Greenwood's formula for variance\n        se = survival * np.sqrt(var_sum) if var_sum > 0 else 0\n        ci_lower = max(0, survival - 1.96 * se)\n        ci_upper = min(1, survival + 1.96 * se)\n        results.append({\"time\": t, \"survival\": survival, \"ci_lower\": ci_lower, \"ci_upper\": ci_upper, \"n_at_risk\": n})\n        n_at_risk -= len(at_time)\n\n    return pd.DataFrame(results)\n\n\n# Compute Kaplan-Meier for each group\nkm_treatment = kaplan_meier(treatment[\"time\"], treatment[\"event\"])\nkm_treatment[\"group\"] = \"Treatment\"\nkm_control = kaplan_meier(control[\"time\"], control[\"event\"])\nkm_control[\"group\"] = \"Control\"\nkm_data = pd.concat([km_treatment, km_control], ignore_index=True)\n\n# Find censored observations for tick marks\ncensored_treatment = treatment[treatment[\"event\"] == 0].copy()\ncensored_control = control[control[\"event\"] == 0].copy()\n\n\ndef get_survival_at_time(km_df, t):\n    \"\"\"Get survival probability at a given time.\"\"\"\n    km_df = km_df.sort_values(\"time\")\n    idx = km_df[km_df[\"time\"] <= t].index\n    if len(idx) == 0:\n        return 1.0\n    return km_df.loc[idx[-1], \"survival\"]\n\n\ncensored_treatment[\"survival\"] = censored_treatment[\"time\"].apply(lambda t: get_survival_at_time(km_treatment, t))\ncensored_treatment[\"group\"] = \"Treatment\"\ncensored_control[\"survival\"] = censored_control[\"time\"].apply(lambda t: get_survival_at_time(km_control, t))\ncensored_control[\"group\"] = \"Control\"\ncensored_data = pd.concat([censored_treatment, censored_control], ignore_index=True)\n\n# Plot\nplot = (\n    ggplot()\n    # Confidence interval ribbons\n    + geom_ribbon(aes(x=\"time\", ymin=\"ci_lower\", ymax=\"ci_upper\", fill=\"group\"), data=km_data, alpha=0.2)\n    # Step functions for survival curves\n    + geom_step(aes(x=\"time\", y=\"survival\", color=\"group\"), data=km_data, size=1.5)\n    # Censored observation tick marks\n    + geom_point(\n        aes(x=\"time\", y=\"survival\", color=\"group\"),\n        data=censored_data,\n        shape=3,  # Plus sign for censoring ticks\n        size=4,\n        stroke=2,\n    )\n    + scale_color_manual(values=IMPRINT)\n    + scale_fill_manual(values=IMPRINT)\n    + labs(\n        x=\"Time (months)\",\n        y=\"Survival Probability\",\n        title=\"survival-kaplan-meier · letsplot · pyplots.ai\",\n        color=\"Group\",\n        fill=\"Group\",\n    )\n    + theme_minimal()\n    + theme(\n        plot_background=element_rect(fill=PAGE_BG, color=PAGE_BG),\n        panel_background=element_rect(fill=PAGE_BG),\n        panel_grid_major=element_line(color=INK_MUTED, size=0.3),\n        axis_title=element_text(size=20, color=INK),\n        axis_text=element_text(size=16, color=INK_SOFT),\n        axis_line=element_line(color=INK_SOFT, size=0.5),\n        plot_title=element_text(size=24, color=INK),\n        legend_background=element_rect(fill=ELEVATED_BG, color=INK_SOFT),\n        legend_text=element_text(size=16, color=INK_SOFT),\n        legend_title=element_text(size=18, color=INK),\n    )\n    + ggsize(1600, 900)\n)\n\n# Save\nggsave(plot, f\"plot-{THEME}.png\", scale=3)\nggsave(plot, f\"plot-{THEME}.html\")\n\n# lets-plot creates a subdirectory; move files to current directory\nif os.path.exists(\"lets-plot-images\"):\n    for file in os.listdir(\"lets-plot-images\"):\n        shutil.move(f\"lets-plot-images/{file}\", file)\n    os.rmdir(\"lets-plot-images\")\n"}