Add centered FFT example to fftshift docs (#51223)

Summary:
Closes #51022

Pull Request resolved: #51223

Reviewed By: malfet

Differential Revision: D26110201

Pulled By: mruberry

fbshipit-source-id: c659c5dca30eda4b67ed6d931a93de9a33e72895

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"""