fft#

nvmath.fft.fft(
x,
*,
axes=None,
direction=None,
选项=None,
执行=None,
前言=None,
后记=None,
=None,
检查数据类型=None,
)[源代码]#

fft(operand, axes=None, direction=None, options=None, execution=None, prolog=None, epilog=None, stream=None)

对提供的复数操作数执行 N 维复数到复数 (C2C) FFT。

参数:
  • operand – 张量(类似 ndarray 的对象)。目前支持的类型为 numpy.ndarray, cupy.ndarray, 和 torch.Tensor

  • axes – 执行 FFT 的维度。 axes[-1] 是 rfft 的“最后转换”轴。目前,要求轴是连续的,并且包含第一个或最后一个维度。仅支持高达 3D FFT。

  • options – 将 FFT 的选项指定为 FFTOptions 对象。或者,也可以提供包含 FFTOptions 构造函数参数的 dict。如果未指定,则该值将设置为默认构造的 FFTOptions 对象。

  • prolog – 以 LTO-IR 格式提供设备可调用函数,以用作加载回调,作为 DeviceCallable 类型的对象。或者,也可以提供包含 DeviceCallable 构造函数参数的 dict。默认情况下没有 prolog。目前,回调仅在 CUDA 执行中受支持。

  • epilog – 以 LTO-IR 格式提供设备可调用函数,以用作存储回调,作为 DeviceCallable 类型的对象。或者,也可以提供包含 DeviceCallable 构造函数参数的 dict。默认情况下没有 epilog。目前,回调仅在 CUDA 执行中受支持。

  • stream – 提供用于执行操作的 CUDA 流。可接受的输入包括 cudaStream_t (作为 Python int), cupy.cuda.Stream, 和 torch.cuda.Stream。如果未提供流,将使用操作数包中的当前流。

返回:

转换后的操作数,它保留与输入相同的数据类型和形状。它保留在同一设备上,并使用与输入操作数相同的包。

另请参阅

ifft(), irfft(), rfft(), FFT

示例

>>> import cupy as cp
>>> import nvmath

在 GPU 上创建一个 3D complex128 ndarray

>>> shape = 256, 256, 256
>>> a = cp.random.rand(*shape, dtype=cp.float64) + 1j * cp.random.rand(
...     *shape, dtype=cp.float64
... )

使用 fft() 执行 3D C2C FFT。结果 r 也是一个 CuPy complex128 ndarray

>>> r = nvmath.fft.fft(a)

用户还可以沿维度子集执行 FFT,例如,沿前两个维度执行 2D C2C FFT,沿最后一个维度进行批处理

>>> axes = 0, 1
>>> r = nvmath.fft.fft(a, axes=axes)

对于 C2C 类型 FFT 操作,可以直接就地计算输出,从而覆盖输入操作数。这可以使用 FFT 的选项来指定

>>> o = nvmath.fft.FFTOptions(inplace=True)
>>> r = nvmath.fft.fft(a, options=o)
>>> r is a
True

有关可用选项的完整列表,请参阅 FFTOptions

默认情况下使用包的当前流,但可以显式地为 FFT 操作提供流。如果 FFT 操作数是在不同的流上计算的,则可以这样做,例如

>>> s = cp.cuda.Stream()
>>> with s:
...     a = cp.random.rand(*shape) + 1j * cp.random.rand(*shape)
>>> r = nvmath.fft.fft(a, stream=s)

上述操作在流 s 上运行,并相对于输入计算进行排序。

在 CPU 上创建一个 NumPy ndarray。

>>> import numpy as np
>>> b = np.random.rand(*shape) + 1j * np.random.rand(*shape)

将 NumPy ndarray 提供给 fft(),结果也是一个 NumPy ndarray

>>> r = nvmath.fft.fft(b)

注解

  • 此函数仅接受复数操作数进行 C2C 变换。如果用户希望对实数输入执行完整 FFT 变换,请将输入转换为相应的复数数据类型。

  • 此函数是 FFT 的便捷包装器,专门用于单次使用。可以使用状态 API 通过 FFT.execute() 中的默认 direction 参数执行相同的计算。

更多示例可以在 nvmath/examples/fft 目录中找到。