diff --git a/pyhealth/metrics/regression.py b/pyhealth/metrics/regression.py index 2d0d02ed5..094b21ab1 100644 --- a/pyhealth/metrics/regression.py +++ b/pyhealth/metrics/regression.py @@ -51,11 +51,10 @@ def regression_metrics_fn( output = {} for metric in metrics: if metric == "kl_divergence": - # Work on copies to avoid mutating x/x_rec for subsequent metrics - x_kl = x.copy() - x_rec_kl = x_rec.copy() - x_kl[x_kl < 1e-6] = 1e-6 - x_rec_kl[x_rec_kl < 1e-6] = 1e-6 + # Clip into copies without truncating the floor for integer inputs. + # Keep x/x_rec unchanged for subsequent metrics. + x_kl = np.clip(x, 1e-6, None) + x_rec_kl = np.clip(x_rec, 1e-6, None) x_kl = x_kl / np.sum(x_kl) x_rec_kl = x_rec_kl / np.sum(x_rec_kl) kl_divergence = np.sum(x_rec_kl * np.log(x_rec_kl / x_kl)) diff --git a/tests/core/test_regression_integer_kl.py b/tests/core/test_regression_integer_kl.py new file mode 100644 index 000000000..7dee8b3e4 --- /dev/null +++ b/tests/core/test_regression_integer_kl.py @@ -0,0 +1,86 @@ +"""Regression tests for KL divergence on integer-valued arrays. + +Author: Pema Sangpo Lama +Description: Exercise zero clipping using tiny synthetic count arrays. +""" + +import math +import unittest + +import numpy as np + +from pyhealth.metrics import regression_metrics_fn, timeseries_regression_metrics_fn + + +class TestIntegerKLDivergence(unittest.TestCase): + def setUp(self): + # KL([1, eps] / (1 + eps) || [eps, 1] / (1 + eps)). + eps = 1e-6 + self.expected = (1 - eps) / (1 + eps) * math.log(1 / eps) + + def test_zero_counts_are_clipped_in_either_input(self): + for dtype in (np.int32, np.int64, np.uint64): + for integer_target, integer_prediction in ( + (True, False), + (False, True), + (True, True), + ): + with self.subTest( + dtype=dtype, target=integer_target, prediction=integer_prediction + ): + x = np.array([0, 1], dtype=dtype if integer_target else float) + x_rec = np.array( + [1, 0], dtype=dtype if integer_prediction else float + ) + with np.errstate(divide="raise", invalid="raise"): + result = regression_metrics_fn(x, x_rec, ["kl_divergence"]) + self.assertAlmostEqual(result["kl_divergence"], self.expected) + + def test_all_zero_integer_inputs_give_zero_divergence(self): + for dtype in (np.int32, np.int64, np.uint64): + with self.subTest(dtype=dtype): + with np.errstate(divide="raise", invalid="raise"): + result = regression_metrics_fn( + np.zeros(2, dtype=dtype), + np.zeros(2, dtype=dtype), + ["kl_divergence"], + ) + self.assertEqual(result["kl_divergence"], 0.0) + + def test_floating_point_inputs_still_match_clipped_divergence(self): + for dtype, places in ((np.float32, 5), (np.float64, 7)): + with self.subTest(dtype=dtype): + result = regression_metrics_fn( + np.array([0, 1], dtype=dtype), + np.array([1, 0], dtype=dtype), + ["kl_divergence"], + ) + self.assertAlmostEqual(result["kl_divergence"], self.expected, places) + + def test_other_metrics_and_inputs_keep_original_values(self): + x = np.array([0, 1], dtype=np.int64) + x_rec = np.array([1, 0], dtype=np.int64) + for metrics in (None, ["mse", "mae", "kl_divergence"]): + with self.subTest(metrics=metrics): + with np.errstate(divide="raise", invalid="raise"): + result = regression_metrics_fn(x, x_rec, metrics) + self.assertAlmostEqual(result["kl_divergence"], self.expected) + self.assertEqual(result["mse"], 1.0) + self.assertEqual(result["mae"], 1.0) + np.testing.assert_array_equal(x, [0, 1]) + np.testing.assert_array_equal(x_rec, [1, 0]) + self.assertEqual(x.dtype, np.dtype(np.int64)) + self.assertEqual(x_rec.dtype, np.dtype(np.int64)) + + def test_timeseries_integer_counts_use_safe_clipping(self): + with np.errstate(divide="raise", invalid="raise"): + result = timeseries_regression_metrics_fn( + [np.array([0]), np.array([1])], + [np.array([1]), np.array([0])], + ["kl_divergence"], + ) + self.assertAlmostEqual(result["kl_divergence"], self.expected) + + +if __name__ == "__main__": + unittest.main()