diff --git a/torch/fft/__init__.py b/torch/fft/__init__.py index 7efdb04a52d3..4131403066c8 100644 --- a/torch/fft/__init__.py +++ b/torch/fft/__init__.py @@ -865,6 +865,10 @@ Reorders n-dimensional FFT data, as provided by :func:`~torch.fft.fftn`, to have negative frequency terms first. +This performs a periodic shift of n-dimensional data such that the origin +``(0, ..., 0)`` is moved to the center of the tensor. Specifically, to +``input.shape[dim] // 2`` in each selected dimension. + Note: By convention, the FFT returns positive frequency terms first, followed by the negative frequencies in reverse order, so that ``f[-i]`` for all @@ -889,10 +893,10 @@ >>> f = torch.fft.fftfreq(4) >>> f - tensor([ 0.0000, 0.2500, -0.5000, -0.2500]) + tensor([ 0.0000, 0.2500, -0.5000, -0.2500]) - >>> torch.fftshift(f) - tensor([-0.5000, -0.2500, 0.0000, 0.2500]) + >>> torch.fft.fftshift(f) + tensor([-0.5000, -0.2500, 0.0000, 0.2500]) Also notice that the Nyquist frequency term at ``f[2]`` was moved to the beginning of the tensor. @@ -914,6 +918,28 @@ [-1.9000, -0.9000, 0.1000, 1.1000, 2.1000], [-1.8000, -0.8000, 0.2000, 1.2000, 2.2000]]) + :func:`~torch.fft.fftshift` can also be useful for spatial data. If our + data is defined on a centered grid (``[-(N//2), (N-1)//2]``) then we can + use the standard FFT defined on an uncentered grid (``[0, N)``) by first + applying an :func:`~torch.fft.ifftshift`. + + >>> x_centered = torch.arange(-5, 5) + >>> x_uncentered = torch.fft.ifftshift(x_centered) + >>> fft_uncentered = torch.fft.fft(x_uncentered) + + Similarly, we can convert the frequency domain components to centered + convention by applying :func:`~torch.fft.fftshift`. + + >>> fft_centered = torch.fft.fftshift(fft_uncentered) + + The inverse transform, from centered Fourier space back to centered spatial + data, can be performed by applying the inverse shifts in reverse order: + + >>> x_centered_2 = torch.fft.fftshift(torch.fft.ifft(torch.fft.ifftshift(fft_centered))) + >>> torch.allclose(x_centered.to(torch.complex64), x_centered_2) + True + + """) ifftshift = _add_docstr(_fft.fft_ifftshift, r"""