Merge pull request #1672 from haeunhwangbo/fix-add-at-risk-counts-numpy-2-4

Fix add_at_risk_counts for NumPy>=2.4

Author: CamDavidsonPilon

lifelines/plotting.py

@@ -540,7 +540,10 @@ def add_at_risk_counts(
                     event_table_slice.loc[:tick, ["at_risk", "censored", "observed"]]
                     .agg(
                         {
-                            "at_risk": lambda x: x.tail(1).values,
+                            # `Series.tail(1).values` is a 1D array of length 1. In NumPy>=2.4,
+                            # `int(np.array([1]))` raises `TypeError: only 0-dimensional arrays can be converted to Python scalars`.
+                            # Extract a Python scalar for compatibility.
+                            "at_risk": lambda x: x.tail(1).values.item(),
                             "censored": "sum",
                             "observed": "sum",
                         }
@@ -554,7 +557,7 @@ def add_at_risk_counts(
                     )
                     .fillna(0)
                 )
-                counts.extend([int(c) for c in event_table_slice.loc[rows_to_show]])
+                counts.extend([int(np.asarray(c).item()) for c in event_table_slice.loc[rows_to_show]])
             else:
                 counts.extend([0 for _ in range(n_rows)])
         if n_rows > 1:

lifelines/tests/test_plotting.py

@@ -44,6 +44,23 @@ def waltons():
     return load_waltons()[["T", "E"]].iloc[:50]
 
 
+def test_add_at_risk_counts_is_numpy_scalar_compatible():
+    # Regression test for https://github.com/CamDavidsonPilon/lifelines/issues/1671:
+    # `add_at_risk_counts` previously produced 1D NumPy arrays (length 1) for "At risk" counts,
+    # which breaks `int(c)` on NumPy>=2.4 (`TypeError: only 0-dimensional arrays can be converted to Python scalars`).
+    matplotlib = pytest.importorskip("matplotlib")
+    matplotlib.use("Agg", force=True)
+    from matplotlib import pyplot as plt
+
+    kmf = KaplanMeierFitter().fit(
+        np.random.exponential(10, size=(100,)),
+        np.random.binomial(1, 0.8, size=(100,)),
+    )
+    ax = kmf.plot_survival_function()
+    add_at_risk_counts(kmf, ax=ax)
+    plt.close(ax.figure)
+
+
 @pytest.mark.skipif("DISPLAY" not in os.environ, reason="requires display")
 class TestPlotting:
     @pytest.fixture