diff --git a/pyhealth/metrics/calibration.py b/pyhealth/metrics/calibration.py index 32e27b617..99f241e32 100644 --- a/pyhealth/metrics/calibration.py +++ b/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): diff --git a/tests/core/test_calibration_binary_ece.py b/tests/core/test_calibration_binary_ece.py new file mode 100644 index 000000000..0b28f02b5 --- /dev/null +++ b/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])