Fix CodeRabbit review issues for radial Fourier transform

- Add device parameter to _compute_radial_coordinates to prevent CPU/GPU mismatch
- Fix frequency mask expansion for multi-dimensional inputs
- Add reconstruction accuracy test assertion (with proper magnitude+phase for inverse)
- Add Raises section to docstring
- Remove unused import
- Address all review comments

Author: Hitendrasinhdata7

monai/transforms/signal/radial_fourier.py

@@ -25,10 +25,10 @@
 
 from monai.config import NdarrayOrTensor
 from monai.transforms.transform import Transform
-from monai.utils import convert_data_type, optional_import
+from monai.utils import convert_data_type
 
 # Optional imports for type checking
-spatial, _ = optional_import("monai.utils", name="spatial")
+# spatial, _ = optional_import("monai.utils", name="spatial")  # Commented out unused import
 
 
 class RadialFourier3D(Transform):
@@ -59,6 +59,10 @@ class RadialFourier3D(Transform):
         >>> transform = RadialFourier3D(radial_bins=64, return_magnitude=True)
         >>> image = torch.randn(1, 128, 128, 96)  # Batch, Height, Width, Depth
         >>> result = transform(image)  # Shape: (1, 64)
+
+    Raises:
+        ValueError: If max_frequency not in (0.0, 1.0], radial_bins < 1, or both
+            return_magnitude and return_phase are False.
     """
 
     def __init__(
@@ -89,12 +93,13 @@ def __init__(
         if not return_magnitude and not return_phase:
             raise ValueError("At least one of return_magnitude or return_phase must be True")
 
-    def _compute_radial_coordinates(self, shape: tuple[int, ...]) -> torch.Tensor:
+    def _compute_radial_coordinates(self, shape: tuple[int, ...], device: torch.device = None) -> torch.Tensor:
         """
         Compute radial distance from frequency domain center.
 
         Args:
             shape: spatial dimensions (D, H, W) or (H, W, D) depending on dims order.
+            device: device to create tensor on.
 
         Returns:
             Tensor of same spatial shape with radial distances.
@@ -103,7 +108,7 @@ def _compute_radial_coordinates(self, shape: tuple[int, ...]) -> torch.Tensor:
         coords = []
         for dim_size in shape:
             # Create frequency range from -0.5 to 0.5
-            freq = torch.fft.fftfreq(dim_size)
+            freq = torch.fft.fftfreq(dim_size, device=device)
             coords.append(freq)
 
         # Create meshgrid and compute radial distance
@@ -176,7 +181,7 @@ def __call__(self, img: NdarrayOrTensor) -> NdarrayOrTensor:
             spectrum = spectrum / norm_factor
 
         # Compute radial coordinates
-        radial_coords = self._compute_radial_coordinates(spatial_shape)
+        radial_coords = self._compute_radial_coordinates(spatial_shape, device=spectrum.device)
 
         # Apply radial binning if requested
         if self.radial_bins is not None:
@@ -217,7 +222,8 @@ def __call__(self, img: NdarrayOrTensor) -> NdarrayOrTensor:
             if self.max_frequency < 1.0:
                 freq_mask = radial_coords <= (self.max_frequency * 0.5)
                 # Expand mask to match spectrum dimensions
-                for _ in range(len(self.spatial_dims)):
+                n_non_spatial = len(spectrum.shape) - len(spatial_shape)
+                for _ in range(n_non_spatial):
                     freq_mask = freq_mask.unsqueeze(0)
                 spectrum = spectrum * freq_mask
 

tests/test_radial_fourier.py

@@ -76,7 +76,7 @@ def test_normalization(self):
     def test_inverse_transform(self):
         """Test approximate inverse transform."""
         # Use full spectrum for invertibility
-        transform = RadialFourier3D(radial_bins=None, normalize=True)
+        transform = RadialFourier3D(radial_bins=None, normalize=True, return_magnitude=True, return_phase=True)
 
         # Forward transform
         spectrum = transform(self.test_image_3d)
@@ -87,6 +87,9 @@ def test_inverse_transform(self):
         # Should have same shape
         self.assertEqual(reconstructed.shape, self.test_image_3d.shape)
 
+        # Should approximately reconstruct original
+        self.assertTrue(torch.allclose(reconstructed, self.test_image_3d, atol=1e-5))
+
     def test_deterministic(self):
         """Test that transform is deterministic."""
         transform = RadialFourier3D(radial_bins=32)