Add weight clipping and comments

Quantco · Jul 18, 2024 · 45554cb · 45554cb
1 parent 95cabbb
commit 45554cb
Showing 1 changed file with 6 additions and 1 deletion.
diff --git a/docs/examples/example_survival.ipynb b/docs/examples/example_survival.ipynb
@@ -149,7 +149,7 @@
     "from lifelines import KaplanMeierFitter\n",
     "\n",
     "\n",
-    "def estimate_ipcw_km(outcome, censored):\n",
+    "def estimate_ipcw_km(outcome, censored, clip=0.95):\n",
     "    \"\"\"\n",
     "    Estimate Inverse Probability of Censoring Weights (IPCW) by nonparametric Kaplan-Meier method\n",
     "    (which is valid in case of noninformative censoring)\n",
@@ -164,8 +164,13 @@
     "    inverse_prob_censoring_weight = pd.Series(\n",
     "        index=outcome.index, data=inverse_prob_censoring_weight.values\n",
     "    )\n",
+    "    # We only want to use complete observations, therefore we give a weight of 0 to censored observations\n",
     "    inverse_prob_censoring_weight[censored.astype(bool)] = 0\n",
     "\n",
+    "    # We clip the weights to a percentile to avoid exploding weights\n",
+    "    clip_quantile = inverse_prob_censoring_weight.quantile(clip)\n",
+    "    inverse_prob_censoring_weight = inverse_prob_censoring_weight.clip(0, clip_quantile)\n",
+    "\n",
     "    return inverse_prob_censoring_weight\n",
     "\n",
     "\n",