Fix ECE confidence binary crash

pyhealth/metrics/calibration.py

@@ -137,6 +137,12 @@ def ece_confidence_binary(prob:np.ndarray, label:np.ndarray, bins=20, adaptive=F
 
     Similar to :func:`ece_confidence_multiclass`, but on class 1 instead of the top-prediction.
 
+    Examples:
+        >>> prob = np.array([0.1, 0.8])
+        >>> label = np.array([0, 1])
+        >>> ece = ece_confidence_binary(prob, label, bins=2)
+        >>> 0.0 <= ece <= 1.0
+        True
 
     Args:
         prob (np.ndarray): (N, C)
@@ -147,7 +153,13 @@ def ece_confidence_binary(prob:np.ndarray, label:np.ndarray, bins=20, adaptive=F
             of points. Defaults to False.
     """
 
-    df = pd.DataFrame({'acc': label[:,0], 'conf': prob[:,0]})
+    prob = np.asarray(prob)
+    label = np.asarray(label)
+
+    conf = prob[:, 0 if prob.shape[1] == 1 else 1] if prob.ndim > 1 else prob
+    acc = label[:, 0 if label.shape[1] == 1 else 1] if label.ndim > 1 else label
+
+    df = pd.DataFrame({'acc': acc, 'conf': conf})
     return _ECE_confidence(df, bins, adaptive)[1]
 
 def ece_classwise(prob, label, bins=20, threshold=0., adaptive=False):

tests/core/test_calibration_binary_ece.py

@@ -0,0 +1,47 @@
+import unittest
+from unittest.mock import patch
+
+import numpy as np
+
+from pyhealth.metrics import binary_metrics_fn
+from pyhealth.metrics.calibration import ece_confidence_binary
+
+
+class TestBinaryECE(unittest.TestCase):
+    def test_binary_metrics_fn_ece_does_not_crash(self):
+        y_true = np.array([0, 0, 1, 1, 0, 1])
+        y_prob = np.array([0.1, 0.4, 0.35, 0.8, 0.2, 0.7])
+        for metric in ("ECE", "ECE_adapt"):
+            out = binary_metrics_fn(y_true, y_prob, metrics=[metric])
+            self.assertIn(metric, out)
+            self.assertTrue(np.isfinite(out[metric]))
+            self.assertGreaterEqual(out[metric], 0.0)
+            self.assertLessEqual(out[metric], 1.0)
+
+    def test_two_dim_inputs_use_positive_class(self):
+        prob = np.array([[0.2, 0.8], [0.7, 0.3]])
+        label = np.array([[0, 1], [1, 0]])
+
+        with patch(
+            "pyhealth.metrics.calibration._ECE_confidence",
+            return_value=(None, 0.0),
+        ) as ece:
+            ece_confidence_binary(prob, label)
+
+        frame = ece.call_args.args[0]
+        np.testing.assert_array_equal(frame["conf"].to_numpy(), prob[:, 1])
+        np.testing.assert_array_equal(frame["acc"].to_numpy(), label[:, 1])
+
+    def test_single_column_inputs_use_only_column(self):
+        prob = np.array([[0.2], [0.7]])
+        label = np.array([[0], [1]])
+
+        with patch(
+            "pyhealth.metrics.calibration._ECE_confidence",
+            return_value=(None, 0.0),
+        ) as ece:
+            ece_confidence_binary(prob, label)
+
+        frame = ece.call_args.args[0]
+        np.testing.assert_array_equal(frame["conf"].to_numpy(), prob[:, 0])
+        np.testing.assert_array_equal(frame["acc"].to_numpy(), label[:, 0])