Handle constant B-spline mutual information inputs

monai/losses/image_dissimilarity.py

@@ -214,7 +214,7 @@ def __init__(
                       IEEE Transactions in Medical Imaging. Vol.22, No.1,
                       January 2003. pp.120-128.
 
-            num_bins: number of bins for intensity
+            num_bins: number of bins for intensity. The B-spline kernel requires more than 4 bins.
             sigma_ratio: a hyper param for gaussian function
             reduction: {``"none"``, ``"mean"``, ``"sum"``}
                 Specifies the reduction to apply to the output. Defaults to ``"mean"``.
@@ -224,13 +224,20 @@ def __init__(
                 - ``"sum"``: the output will be summed.
             smooth_nr: a small constant added to the numerator to avoid nan.
             smooth_dr: a small constant added to the denominator to avoid nan.
+
+        Raises:
+            ValueError: if ``num_bins`` is not positive, or if a B-spline
+                kernel is configured with four or fewer bins.
         """
         super().__init__(reduction=LossReduction(reduction).value)
         if num_bins <= 0:
             raise ValueError(f"num_bins must > 0, got {num_bins}")
         bin_centers = torch.linspace(0.0, 1.0, num_bins)  # (num_bins,)
         sigma = torch.mean(bin_centers[1:] - bin_centers[:-1]) * sigma_ratio
         self.kernel_type = look_up_option(kernel_type, ["gaussian", "b-spline"])
+        # B-spline windowing reserves two padding bins at each boundary.
+        if self.kernel_type == "b-spline" and num_bins <= 4:
+            raise ValueError(f"num_bins must be greater than 4 for b-spline kernel, got {num_bins}")
         self.num_bins = num_bins
         # declared as buffers so they move with the module (e.g. ``.to(device)``); only populated for the
         # gaussian kernel, hence the ``Tensor`` annotation reflects the type at the use sites in that path.
@@ -281,13 +288,20 @@ def parzen_windowing_b_spline(self, img: torch.Tensor, order: int) -> tuple[torc
         # Note that there can still be non-zero bin values in the padded region,
         # it's just that these bins will never be a central bin for the Parzen
         # window.
-        _max, _min = torch.max(img), torch.min(img)
+        compute_img = img.float() if img.dtype in (torch.float16, torch.bfloat16) else img
+        _max, _min = torch.max(compute_img), torch.min(compute_img)
         padding = 2
-        bin_size = (_max - _min) / (self.num_bins - 2 * padding)
+        interior_bins = self.num_bins - 2 * padding
+        value_range = _max - _min
+        raw_bin_size = value_range / interior_bins
+        # Smaller bins either underflow the compute dtype or require a
+        # normalization slope that cannot be represented by the input dtype.
+        minimum_bin_size = max(torch.finfo(compute_img.dtype).tiny, 1.0 / torch.finfo(img.dtype).max)
+        bin_size = torch.where(raw_bin_size >= minimum_bin_size, raw_bin_size, torch.ones_like(raw_bin_size))
         norm_min = torch.div(_min, bin_size) - padding
 
         # assign bin/window index to each voxel
-        window_term = torch.div(img, bin_size) - norm_min  # B[NDHW]
+        window_term = torch.div(compute_img, bin_size) - norm_min  # B[NDHW]
         # make sure the extreme values are in valid (non-padded) bins
         window_term = torch.clamp(window_term, padding, self.num_bins - padding - 1)  # B[NDHW]
         window_term = window_term.reshape(window_term.shape[0], -1, 1)  # (batch, num_sample, 1)

tests/losses/image_dissimilarity/test_global_mutual_information_loss.py

@@ -135,6 +135,11 @@ def test_b_spline_bin_centers_exists_as_none(self):
 
         self.assertIsNone(loss.bin_centers)
 
+    def test_b_spline_num_bins_must_allow_padding(self):
+        """Verify B-spline bin counts leave room for boundary padding."""
+        with self.assertRaisesRegex(ValueError, "num_bins must be greater than 4"):
+            GlobalMutualInformationLoss(kernel_type="b-spline", num_bins=4)
+
     @parameterized.expand(
         [
             (torch.ones((1, 2), dtype=torch.float), torch.ones((1, 3), dtype=torch.float)),  # mismatched_simple_dims
@@ -164,6 +169,65 @@ def test_ill_opts(self, num_bins, reduction, expected_exception, expected_messag
             GlobalMutualInformationLoss(num_bins=num_bins, reduction=reduction)(pred, target)
 
 
+class TestGlobalMutualInformationLossBSpline(unittest.TestCase):
+    """Test B-spline mutual information on degenerate intensity ranges."""
+
+    @parameterized.expand(["prediction", "target"])
+    def test_b_spline_single_constant_input_is_finite(self, constant_input):
+        """Verify either independently constant input yields finite gradients."""
+        varying = torch.linspace(0.0, 1.0, 64).reshape(1, 1, 8, 8)
+        if constant_input == "prediction":
+            pred = torch.zeros_like(varying, requires_grad=True)
+            target = varying
+        else:
+            pred = varying.clone().requires_grad_()
+            target = torch.ones_like(varying)
+        loss = GlobalMutualInformationLoss(kernel_type="b-spline")
+
+        result = loss(pred, target)
+
+        self.assertTrue(torch.isfinite(result))
+        result.backward()
+        self.assertIsNotNone(pred.grad)
+        self.assertTrue(torch.isfinite(pred.grad).all())
+
+    def test_b_spline_constant_half_precision_images_are_finite(self):
+        """Verify constant float16 inputs survive overflow-sized bin distances."""
+        pred = torch.zeros((1, 1, 8, 8), dtype=torch.float16, requires_grad=True)
+        target = torch.ones_like(pred)
+        loss = GlobalMutualInformationLoss(kernel_type="b-spline", num_bins=64)
+
+        result = loss(pred, target)
+
+        self.assertTrue(torch.isfinite(result))
+        self.assertAlmostEqual(result.item(), 0.0, places=6)
+        result.backward()
+        self.assertIsNotNone(pred.grad)
+        self.assertTrue(torch.isfinite(pred.grad).all())
+
+    @parameterized.expand(
+        [
+            ("float16_tiny", torch.float16, torch.finfo(torch.float16).tiny / 2),
+            ("bfloat16_tiny", torch.bfloat16, torch.finfo(torch.bfloat16).tiny / 2),
+            ("float32_tiny", torch.float32, torch.finfo(torch.float32).tiny / 2),
+            ("float16_large", torch.float16, 65000.0),
+        ]
+    )
+    def test_b_spline_nonzero_ranges_are_finite(self, _, dtype, maximum):
+        """Verify extreme nonzero ranges yield finite loss and gradients."""
+        values = torch.tensor([0.0, maximum, maximum, 0.0], dtype=dtype).reshape(1, 1, 2, 2)
+        pred = values.clone().requires_grad_()
+        target = torch.flip(values, dims=(-1,))
+        loss = GlobalMutualInformationLoss(kernel_type="b-spline", num_bins=64)
+
+        result = loss(pred, target)
+
+        self.assertTrue(torch.isfinite(result))
+        result.backward()
+        self.assertIsNotNone(pred.grad)
+        self.assertTrue(torch.isfinite(pred.grad).all())
+
+
 class TestGlobalMutualInformationLossBuffers(unittest.TestCase):
     def test_gaussian_kernel_registers_buffers(self):
         """Verify gaussian kernel registers preterm and bin_centers as non-trainable, non-persistent buffers."""