Handle constant B-spline mutual information inputs

monai/losses/image_dissimilarity.py

@@ -214,7 +214,8 @@ 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 Gaussian kernel requires
+                more than 1 bin, and 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 +225,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 a Gaussian kernel is configured with one or fewer
+                bins, 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"])
+        if self.kernel_type == "gaussian" and num_bins <= 1:
+            raise ValueError(f"num_bins must be greater than 1 for gaussian kernel, got {num_bins}")
+        # 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.
@@ -239,6 +247,8 @@ def __init__(
         self.register_buffer("preterm", None, persistent=False)
         self.register_buffer("bin_centers", None, persistent=False)
         if self.kernel_type == "gaussian":
+            bin_centers = torch.linspace(0.0, 1.0, num_bins)  # (num_bins,)
+            sigma = torch.mean(bin_centers[1:] - bin_centers[:-1]) * sigma_ratio
             self.register_buffer("preterm", 1 / (2 * sigma**2), persistent=False)
             self.register_buffer("bin_centers", bin_centers[None, None, ...], persistent=False)
         self.smooth_nr = float(smooth_nr)
@@ -281,13 +291,22 @@ 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 not torch.is_floating_point(img) or 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
+        # Fall back only when the promoted bin size is non-finite or too small
+        # to remain a nonzero normal value in the compute dtype.
+        valid_bin_size = torch.isfinite(raw_bin_size) & (raw_bin_size >= torch.finfo(compute_img.dtype).tiny)
+        bin_size = torch.where(valid_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,16 @@ 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)
+
+    def test_gaussian_num_bins_must_allow_spacing(self):
+        """Verify Gaussian bin counts allow a finite bin-centre spacing."""
+        with self.assertRaisesRegex(ValueError, "num_bins must be greater than 1"):
+            GlobalMutualInformationLoss(kernel_type="gaussian", num_bins=1)
+
     @parameterized.expand(
         [
             (torch.ones((1, 2), dtype=torch.float), torch.ones((1, 3), dtype=torch.float)),  # mismatched_simple_dims
@@ -164,6 +174,104 @@ 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.
+
+        Args:
+            constant_input: Which input (``"prediction"`` or ``"target"``)
+                is held constant.
+        """
+        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())
+
+    def test_b_spline_integer_target_is_supported(self):
+        """Verify integer targets do not fail floating-point range validation."""
+        pred = torch.linspace(0.0, 1.0, 64).reshape(1, 1, 8, 8).requires_grad_()
+        target = torch.arange(64, dtype=torch.int64).reshape(1, 1, 8, 8)
+        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_float16_small_range_preserves_signal(self):
+        """Verify a small nonzero float16 range retains loss and gradient signal."""
+        values = torch.linspace(0.0, 5e-4, 64, dtype=torch.float16).reshape(1, 1, 8, 8)
+        pred = values.clone().requires_grad_()
+        loss = GlobalMutualInformationLoss(kernel_type="b-spline", num_bins=64)
+
+        result = loss(pred, values)
+
+        self.assertTrue(torch.isfinite(result))
+        self.assertLess(result.item(), -1.0)
+        result.backward()
+        self.assertIsNotNone(pred.grad)
+        self.assertTrue(torch.isfinite(pred.grad).all())
+        self.assertGreater(torch.count_nonzero(pred.grad).item(), 0)
+
+    @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, case_name, dtype, maximum):
+        """Verify extreme nonzero ranges yield finite loss and gradients.
+
+        Args:
+            case_name: Descriptive label for the parameterized range case.
+            dtype: Tensor dtype used for the prediction and target.
+            maximum: Nonzero upper endpoint of the tested intensity range.
+        """
+        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."""