Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Add centered FFT example to fftshift docs #51223

Closed
wants to merge 1 commit into from
Closed
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Jump to
Jump to file
Failed to load files.
Diff view
Diff view
32 changes: 29 additions & 3 deletions torch/fft/__init__.py
Expand Up @@ -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
Expand 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])
Copy link
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Found these discrepancies from running the doctests manually. So, it seems fft doctests aren't being run in CI.

Copy link
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Nice fix.


Also notice that the Nyquist frequency term at ``f[2]`` was moved to the
beginning of the tensor.
Expand All @@ -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"""
Expand Down