abtem.core.fft

abtem.core.fft#

Module for handling Fourier transforms and convolution in abTEM.

Module Contents#

Classes#

CachedFFTWConvolution

Functions#

fft2

Compute the 2-dimensional discrete Fourier Transform. Using the FFT library specified in the configuration.

fft2_convolve

Compute the 2-dimensional convolution of an array with a kernel.

fft_crop

Crop an array. It is assumed that the array is centered in Fourier space, this is used for real-space interpolation.

fft_interpolate

Interpolate an array using Fourier space interpolation.

fft_interpolation_masks

Create boolean masks for interpolating between two arrays using Fourier space interpolation.

fft_shift

Shift an array in real space using Fourier space interpolation.

fft_shift_kernel

Create an array representing one or more phase ramp(s) for shifting another array.

fftn

Compute the n-dimensional discrete Fourier Transform. Using the FFT library specified in the configuration.

get_fftw_object

Get a pyfftw object for a given array and a given FFT function. The object is cached and reused if the array shape is the same.

ifft2

Compute the 2-dimensional inverse discrete Fourier Transform. Using the FFT library specified in the configuration.

ifftn

Compute the n-dimensional inverse discrete Fourier Transform. Using the FFT library specified in the configuration.

is_fast_fft_size

Whether an FFT of length n runs on fast radix kernels.

next_fast_fft_size

The smallest length >= n whose prime factors are all in {2, 3, 5, 7}.

warn_if_slow_gpu_fft

Emit the slow-FFT diagnostic for a transform run outside _fft_dispatch.

Data#

U

API#

class abtem.core.fft.CachedFFTWConvolution[source]#

Initialization

abtem.core.fft.U#

‘TypeVar(…)’

abtem.core.fft.fft2(...) → abtem.core.fft.U[source]#

Compute the 2-dimensional discrete Fourier Transform. Using the FFT library specified in the configuration.

abtem.core.fft.fft2_convolve(...) → abtem.core.fft.U[source]#

Compute the 2-dimensional convolution of an array with a kernel.

Parameters:

  • x (ndarray or da.core.Array) – Array to convolve.

  • kernel (ndarray) – Convolution kernel.

  • overwrite_x (bool, optional) – Overwrite the input array.

Returns:

Convolved array.

Return type:

ndarray or da.core.Array

abtem.core.fft.fft_crop(...)[source]#

Crop an array. It is assumed that the array is centered in Fourier space, this is used for real-space interpolation.

Parameters:

  • array (ndarray) – Array to crop.

  • new_shape (tuple of int) – New shape of the array. If the new shape is smaller than the input array, each preceding dimension is treated as a batch dimension.

  • normalize (bool, optional) – If True, renormalize the array to conserve the total amplitude.

Returns:

Cropped array.

Return type:

ndarray

abtem.core.fft.fft_interpolate(...)[source]#

Interpolate an array using Fourier space interpolation.

Parameters:

  • array (ndarray) – Array to interpolate.

  • new_shape (tuple of int) – New shape of the array.

  • normalization (str, optional) – Normalization to apply to the array. Can be ‘values’ or ‘amplitude’.

  • overwrite_x (bool, optional) – Overwrite the input array.

Returns:

Interpolated array.

Return type:

ndarray

abtem.core.fft.fft_interpolation_masks(...) → tuple[ndarray, ndarray][source]#

Create boolean masks for interpolating between two arrays using Fourier space interpolation.

Parameters:

  • shape_in (tuple of int) – Shape of the input array to interpolate from.

  • shape_out (tuple of int) – Shape of the output array to interpolate to.

Returns:

Masks for the input and output arrays.

Return type:

tuple of ndarray

abtem.core.fft.fft_shift(...) → ndarray[source]#

Shift an array in real space using Fourier space interpolation.

Parameters:

  • array (ndarray) – Array to shift.

  • positions (ndarray) – Array of positions to shift the array to. The last dimension should be the number of dimensions to shift.

Returns:

Shifted array

Return type:

ndarray

abtem.core.fft.fft_shift_kernel(...) → ndarray[source]#

Create an array representing one or more phase ramp(s) for shifting another array.

Parameters:

  • positions (ndarray) – Array of positions to shift the array to. The last dimension should be the number of dimensions to shift.

  • shape (tuple) – Shape of the array to shift.

Returns:

Array representing the phase ramp(s).

Return type:

ndarray

abtem.core.fft.fftn(...) → abtem.core.fft.U[source]#

Compute the n-dimensional discrete Fourier Transform. Using the FFT library specified in the configuration.

abtem.core.fft.get_fftw_object(...)[source]#

Get a pyfftw object for a given array and a given FFT function. The object is cached and reused if the array shape is the same.

Parameters:

  • array (ndarray) – Array to create the FFT object for.

  • name (str) – Name of the FFT function.

  • allow_new_wisdom (bool, optional) – Allow new wisdom to be created.

  • overwrite_x (bool, optional) – Allow the input array to be overwritten.

Returns:

FFTW object.

Return type:

pyfftw.FFTW

abtem.core.fft.ifft2(...) → abtem.core.fft.U[source]#

Compute the 2-dimensional inverse discrete Fourier Transform. Using the FFT library specified in the configuration.

abtem.core.fft.ifftn(...) → abtem.core.fft.U[source]#

Compute the n-dimensional inverse discrete Fourier Transform. Using the FFT library specified in the configuration.

abtem.core.fft.is_fast_fft_size(...) → bool[source]#

Whether an FFT of length n runs on fast radix kernels.

FFT libraries only ship optimized kernels for lengths whose prime factors are small (2, 3, 5 and 7 are supported everywhere). A length with a larger prime factor triggers a generic fallback – on cuFFT the Bluestein algorithm, which pads internally to a power of two, costing several times the arithmetic and, on GPU, a workspace of several times the transform size.

Parameters:

n (int) – The transform length.

Returns:

True if n factorizes completely into 2, 3, 5 and 7.

Return type:

bool

abtem.core.fft.next_fast_fft_size(...) → int[source]#

The smallest length >= n whose prime factors are all in {2, 3, 5, 7}.

Useful for choosing grid sizes (gpts) that avoid the slow large-workspace Bluestein fallback on GPU. Stricter than scipy.fft.next_fast_len, which returns 11-smooth lengths (fast for pocketfft on CPU, but off cuFFT’s documented fast path).

Parameters:

n (int) – The minimum transform length.

Returns:

The next fast transform length.

Return type:

int

abtem.core.fft.warn_if_slow_gpu_fft(...)[source]#

Emit the slow-FFT diagnostic for a transform run outside _fft_dispatch.

A few transforms call xp.fft directly rather than through the wrappers in this module – structure_factor_to_potential does, because the FFTW backend here only ever transforms the trailing two axes and would silently turn its 3D transform into a 2D one. They still deserve the diagnostic, so they can call this alongside. Only GPU arrays are considered: the Bluestein workspace this warns about is a cuFFT concern.