Skip to content

qttools.fft.ffts#

[docs] module qttools.fft.ffts 

  1
  2
  3
  4
  5
  6
  7
  8
  9
 10
 11
 12
 13
 14
 15
 16
 17
 18
 19
 20
 21
 22
 23
 24
 25
 26
 27
 28
 29
 30
 31
 32
 33
 34
 35
 36
 37
 38
 39
 40
 41
 42
 43
 44
 45
 46
 47
 48
 49
 50
 51
 52
 53
 54
 55
 56
 57
 58
 59
 60
 61
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
# Copyright (c) 2024-2026 ETH Zurich and the authors of the qttools package.

from qttools import NDArray, xp

if xp.__name__ == "cupy":
    cache = xp.fft.config.get_plan_cache()


def fft_convolve(a: NDArray, b: NDArray) -> NDArray:
    """Computes the convolution of two arrays using FFT over the first axis (energy axis).

    Parameters
    ----------
    a : NDArray
        First array.
    b : NDArray
        Second array.

    Returns
    -------
    NDArray
        The convolution of the two arrays.
    """
    ne = (
        a.shape[0] + b.shape[0]
    )  # Should not have -1 here (otherwise hilbert transform fails)
    a_fft = xp.fft.fft(a, ne, axis=0)
    b_fft = xp.fft.fft(b, ne, axis=0)
    return xp.fft.ifft(a_fft * b_fft, axis=0)


def fft_circular_convolve(a: NDArray, b: NDArray, axes: tuple[int]) -> NDArray:
    """Computes the circular convolution of two arrays using the FFT.

    Parameters
    ----------
    a : NDArray
        First array.
    b : NDArray
        Second array.
    axes : tuple[int]
        The axes over which to perform the convolution.

    Returns
    -------
    NDArray
        The circular convolution of the two arrays.
    """
    # Extract the shapes of the arrays along the axes as tuples.
    nka = tuple(a.shape[i] for i in axes)
    nkb = tuple(b.shape[i] for i in axes)
    a_fft = xp.fft.fftn(a, nka, axes=axes)
    b_fft = xp.fft.fftn(b, nkb, axes=axes)
    return xp.fft.ifftn(a_fft * b_fft, axes=axes)


def fft_convolve_kpoints(a: NDArray, b: NDArray) -> NDArray:
    """Computes the convolution of two arrays using the FFT.

    The first axis is assumed to be the energy axis, the other
    axes are k-points and the last axis is the orbital index.

    Over the k-point axes, a circular convolution is performed.

    Parameters
    ----------
    a : NDArray
        First array.
    b : NDArray
        Second array.
    Returns
    -------
    NDArray
        The convolution of the two arrays.
    """
    ne = a.shape[0] + b.shape[0] - 1
    nka = a.shape[1:-1]
    nkb = b.shape[1:-1]
    a_fft = xp.fft.fftn(a, (ne,) + nka, axes=(0,) + tuple(range(1, len(nka) + 1)))
    b_fft = xp.fft.fftn(b, (ne,) + nkb, axes=(0,) + tuple(range(1, len(nkb) + 1)))
    return xp.fft.ifftn(a_fft * b_fft, axes=(0,) + tuple(range(1, len(nka) + 1)))


def fft_correlate_kpoints(a: NDArray, b: NDArray) -> NDArray:
    """Computes the correlation of two arrays using FFT.

    The first axis is assumed to be the energy axis, the other
    axes are k-points and the last axis is the orbital index.

    Over the k-point axes, a circular correlation is performed.

    Parameters
    ----------
    a : NDArray
        First array.
    b : NDArray
        Second array.

    Returns
    -------
    NDArray
        The cross-correlation of the two arrays.

    """
    ne = a.shape[0] + b.shape[0] - 1
    nka = a.shape[1:-1]
    nkb = b.shape[1:-1]
    a_fft = xp.fft.fftn(a, (ne,) + nka, axes=(0,) + tuple(range(1, len(nka) + 1)))
    b_fft = xp.fft.fftn(
        xp.flip(b, axis=(0,) + tuple(range(1, len(nkb) + 1))),
        (ne,) + nkb,
        axes=(0,) + tuple(range(1, len(nkb) + 1)),
    )
    # I have to roll the result by one, to get the gamma point at index 0.
    # The ordering is then: [0, 1, 2, ..., -2, -1] in terms of shifts.
    return xp.roll(
        xp.fft.ifftn(a_fft * b_fft, axes=(0,) + tuple(range(1, len(nka) + 1))),
        shift=1,
        axis=tuple(range(1, len(nka) + 1)),
    )