Add PyTorch 1.7 FFT compatibility (#11)

Author: mmuckley

torchkbnufft/nufft/fft_compatibility.py

@@ -0,0 +1,39 @@
+import torch
+from torch import Tensor
+from packaging import version
+
+if version.parse(torch.__version__) >= version.parse("1.7.0"):
+    import torch.fft  # type: ignore
+
+
+def fft_old(image: Tensor, ndim: int, normalized: bool = False) -> Tensor:
+    return torch.fft(image, ndim, normalized)
+
+
+def ifft_old(image: Tensor, ndim: int, normalized: bool = False) -> Tensor:
+    return torch.ifft(image, ndim, normalized)
+
+
+def fft_new(image: Tensor, ndim: int, normalized: bool = False) -> Tensor:
+    norm = "ortho" if normalized else None
+    dims = tuple(range(-ndim, 0))
+
+    image = torch.view_as_real(
+        torch.fft.fftn(  # type: ignore
+            torch.view_as_complex(image.contiguous()), dim=dims, norm=norm
+        )
+    )
+
+    return image
+
+
+def ifft_new(image: Tensor, ndim: int, normalized: bool = False) -> Tensor:
+    norm = "ortho" if normalized else None
+    dims = tuple(range(-ndim, 0))
+    image = torch.view_as_real(
+        torch.fft.ifftn(  # type: ignore
+            torch.view_as_complex(image.contiguous()), dim=dims, norm=norm
+        )
+    )
+
+    return image

torchkbnufft/nufft/fft_functions.py

@@ -1,9 +1,17 @@
 import numpy as np
 import torch
 import torch.nn.functional as F
+from packaging import version
 
 from ..math import complex_mult, conj_complex_mult
 
+if version.parse(torch.__version__) >= version.parse("1.7.0"):
+    from .fft_compatibility import fft_new as fft_fn
+    from .fft_compatibility import ifft_new as ifft_fn
+else:
+    from .fft_compatibility import fft_old as fft_fn
+    from .fft_compatibility import ifft_old as ifft_fn
+
 
 def scale_and_fft_on_image_volume(x, scaling_coef, grid_size, im_size, norm):
     """Applies the FFT and any relevant scaling factors to x.
@@ -43,7 +51,7 @@ def scale_and_fft_on_image_volume(x, scaling_coef, grid_size, im_size, norm):
     # zero pad and fft
     x = F.pad(x, pad_sizes)
     x = x.permute(permute_dims)
-    x = torch.fft(x, grid_size.numel())
+    x = fft_fn(x, grid_size.numel())
     if norm == "ortho":
         x = x / torch.sqrt(torch.prod(grid_size))
     x = x.permute(inv_permute_dims)
@@ -78,7 +86,7 @@ def ifft_and_scale_on_gridded_data(x, scaling_coef, grid_size, im_size, norm):
 
     # do the inverse fft
     x = x.permute(permute_dims)
-    x = torch.ifft(x, grid_size.numel())
+    x = ifft_fn(x, grid_size.numel())
     x = x.permute(inv_permute_dims)
 
     # crop to output size
@@ -140,7 +148,7 @@ def fft_filter(x, kern, norm=None):
     # zero pad and fft
     x = F.pad(x, pad_sizes)
     x = x.permute(permute_dims)
-    x = torch.fft(x, grid_size.numel())
+    x = fft_fn(x, grid_size.numel())
     if norm == "ortho":
         x = x / torch.sqrt(torch.prod(grid_size.to(torch.double)))
     x = x.permute(inv_permute_dims)
@@ -150,7 +158,7 @@ def fft_filter(x, kern, norm=None):
 
     # inverse fft
     x = x.permute(permute_dims)
-    x = torch.ifft(x, grid_size.numel())
+    x = ifft_fn(x, grid_size.numel())
     x = x.permute(inv_permute_dims)
 
     # crop to input size

torchkbnufft/nufft/toep_functions.py

@@ -3,8 +3,12 @@
 
 import numpy as np
 import torch
+from packaging import version
 
-from ..math import absolute
+if version.parse(torch.__version__) >= version.parse("1.7.0"):
+    from .fft_compatibility import fft_new as fft_fn
+else:
+    from .fft_compatibility import fft_old as fft_fn
 
 
 def calc_toep_kernel(adj_ob, om, weights=None):
@@ -117,9 +121,7 @@ def _get_kern(om, weights, flip_list, base_flip, adj_ob):
     inv_permute_dims = (0, 1, kern.ndim - 1) + tuple(range(2, kern.ndim - 1))
 
     # put the kernel in fft space
-    kern = torch.fft(kern.permute(permute_dims), kern.ndim - 3).permute(
-        inv_permute_dims
-    )
+    kern = fft_fn(kern.permute(permute_dims), kern.ndim - 3).permute(inv_permute_dims)
 
     if adj_ob.norm == "ortho":
         kern = kern / torch.sqrt(torch.prod(torch.tensor(kern.shape[3:], dtype=dtype)))