Count infinite histogram samples once in edge buckets

nvflare/app_common/statistics/numpy_utils.py

@@ -69,22 +69,12 @@ def get_std_histogram_buckets(nums: np.ndarray, num_bins: int = 10, br: Optional
         if bucket_count == 0 and num_neginf > 0:
             bucket_low_value = float("-inf")
             bucket_sample_count += num_neginf
-        elif bucket_count == len(counts) - 1 and num_posinf > 0:
+        if bucket_count == len(counts) - 1 and num_posinf > 0:
             bucket_high_value = float("inf")
             bucket_sample_count += num_posinf
 
         histogram_buckets.append(
             Bin(low_value=bucket_low_value, high_value=bucket_high_value, sample_count=bucket_sample_count)
         )
 
-    if buckets is not None and len(buckets) > 0:
-        bucket = None
-        if num_neginf:
-            bucket = Bin(low_value=float("-inf"), high_value=float("-inf"), sample_count=num_neginf)
-        if num_posinf:
-            bucket = Bin(low_value=float("inf"), high_value=float("inf"), sample_count=num_posinf)
-
-        if bucket:
-            histogram_buckets.append(bucket)
-
     return histogram_buckets

tests/unit_test/app_common/statistics/numpy_utils_test.py

@@ -16,8 +16,10 @@
 import pandas as pd
 import pytest
 
-from nvflare.app_common.abstract.statistics_spec import DataType
-from nvflare.app_common.statistics.numpy_utils import dtype_to_data_type
+from nvflare.app_common.abstract.statistics_spec import BinRange, DataType
+from nvflare.app_common.statistics.numeric_stats import accumulate_hists
+from nvflare.app_common.statistics.numpy_utils import dtype_to_data_type, get_std_histogram_buckets
+from nvflare.app_opt.statistics.df.df_core_statistics import DFStatisticsCore
 
 
 class TestDtypeToDataType:
@@ -61,3 +63,58 @@ def test_pandas_nullable_float_dtype(self):
     def test_pandas_nullable_bool_dtype(self):
         # pd.BooleanDtype is a nullable ExtensionDtype — must map to INT, not STRING
         assert dtype_to_data_type(pd.BooleanDtype()) == DataType.INT
+
+
+class TestHistogramInfinityCounts:
+    @pytest.mark.parametrize(
+        "values,num_bins",
+        [
+            ([-1.0, 0.0, 1.0], 3),
+            ([-np.inf, -np.inf, 0.0], 3),
+            ([0.0, np.inf, np.inf], 3),
+            ([-np.inf, 0.0, np.inf], 3),
+            ([-np.inf, 0.0, np.inf], 1),
+            ([-np.inf, np.inf], 1),
+            ([np.nan, -np.inf, 0.0, np.inf], 3),
+            ([-2.0, -1.0, 0.0, 1.0, 2.0], 3),
+            ([], 3),
+        ],
+    )
+    def test_infinities_counted_once_in_edge_buckets(self, values, num_bins):
+        # A read-only strided view verifies that input data is not modified.
+        backing = np.repeat(np.asarray(values, dtype=np.float64), 2)
+        nums = backing[::2]
+        nums.setflags(write=False)
+        counts, edges = np.histogram(nums[np.isfinite(nums)], bins=num_bins, range=(-1.0, 1.0))
+        counts[0] += np.isneginf(nums).sum()
+        counts[-1] += np.isposinf(nums).sum()
+        expected_edges = edges.copy()
+        if np.isneginf(nums).any():
+            expected_edges[0] = -np.inf
+        if np.isposinf(nums).any():
+            expected_edges[-1] = np.inf
+
+        buckets = get_std_histogram_buckets(nums, num_bins, BinRange(-1.0, 1.0))
+        assert len(buckets) == num_bins
+        np.testing.assert_array_equal([b.sample_count for b in buckets], counts)
+        np.testing.assert_array_equal([b.low_value for b in buckets], expected_edges[:-1])
+        np.testing.assert_array_equal([b.high_value for b in buckets], expected_edges[1:])
+        assert sum(b.sample_count for b in buckets) == sum(counts)
+        np.testing.assert_array_equal(nums, values)
+
+    def test_dataframe_and_global_histogram_preserve_sample_count(self):
+        global_histograms = {}
+        expected_count = 0
+        for values in ([-np.inf, 0.0, np.inf], [-np.inf, -0.5, 0.5, np.inf, np.nan]):
+            statistics = DFStatisticsCore()
+            statistics.data = {"train": pd.DataFrame({"value": values})}
+            original = statistics.data["train"].copy(deep=True)
+            histogram = statistics.histogram("train", "value", 3, -1.0, 1.0)
+            expected_count += statistics.count("train", "value")
+            global_histograms = accumulate_hists({"train": {"value": histogram}}, global_histograms)
+            pd.testing.assert_frame_equal(statistics.data["train"], original)
+
+        buckets = global_histograms["train"]["value"].bins
+        assert len(buckets) == 3
+        assert sum(b.sample_count for b in buckets) == expected_count == 7
+        assert [b.sample_count for b in buckets] == [3, 1, 3]