Ctrl+K

Fast Fourier Transform (FFT) Tools#

The following are utility functions for the implementation of various functionals via fast Fourier transform (FFT) related operations. For example, the gradient terms used in semi-local kinetic functionals, such as the Luo-Karasiev-Trickey (LKT) functional, are implemented using FFTs. The following code block shows how these FFT tools can be used to implement the LKT functional, given by

\[T_\text{LKT}[n] = T_\text{vW} + \int_\Omega d^3\mathbf{r} \frac{\tau_\text{TF}(\mathbf{r})}{\cosh(1.3 s)}\]

where \(T_\text{vW}\) is the von Weizsaecker functional, \(\tau_\text{TF}(\mathbf{r})\) is the Thomas-Fermi kinetic energy density and \(s\) is the reduced gradient.

def LuoKarasievTrickey(box_vecs, den):
    TF_ked = 0.3 * (3 * np.pi * np.pi)**(2 / 3) * den.pow(5 / 3)
    kx, ky, kz, k2 = wavevecs(box_vecs, den.shape)
    s = reduced_gradient(kx, ky, kz, den)
    # clamp to avoid s from growing too large, which can cause F_pauli -> 0 and get nan derivatives
    F_pauli = 1 / torch.cosh(1.3 * s.clamp(max=100))
    pauli_T = torch.mean(TF_ked * F_pauli) * torch.abs(torch.linalg.det(box_vecs))
    return Weizsaecker(box_vecs, den) + pauli_T

FFTs are also useful for functionals that involve real space convolutions. Besides the non-local kinetic functionals, the more widely known Hartree functional given by

\[U_\text{Hartree}[n] = \frac{1}{2} \int d^3\mathbf{r} d^3\mathbf{r}'~ \frac{n(\mathbf{r})n(\mathbf{r}')}{|\mathbf{r}-\mathbf{r}'|}\]

also involves convolutions, and hence benefits from an FFT-based implementation motivated by the convolution theorem.

def Hartree(box_vecs, den):
    den_ft = torch.fft.rfftn(den)
    kx, ky, kz, k2 = wavevecs(box_vecs, den.shape)
    coloumb_ft = torch.zeros(k2.shape, dtype=torch.double, device=den.device)
    # set k=0 component to zero. appropriate if the density integrates to
    # zero over the box (e.g. if neutralized by a uniform background charge).
    coloumb_ft[k2 != 0] = 4 * np.pi / k2[k2 != 0]
    pot = torch.fft.irfftn(den_ft * coloumb_ft, den.shape)
    return 0.5 * torch.mean(den * pot) * torch.abs(torch.linalg.det(box_vecs))

Wavevector Tools#

Wavevectors#

functional_tools.wavevecs(box_vecs, shape)#

Generates wavevectors

This is a utility function that generates the wavevectors for a given lattice such that the wavevectors are differentiable with respect to the lattice vectors.

Parameters:

  • box_vecs (torch.Tensor) – Lattice vectors

  • shape (torch.Size or iterable) – Real-space grid shape

Returns:

Wavevectors consistent with real FFTs in the order \(k_x,~k_y,~k_z,~k^2\)

Return type:

torch.Tensor

Gradient#

functional_tools.grad_i(ki, f)#

Computes gradient component

This is a utility function that computes the gradient component or partial spatial derivative

\[\nabla_i f = \frac{\partial f}{\partial r_i}\]

where \(f\) is a given function and \(r_i \in \{x,y,z\}\).

Parameters:

  • ki (torch.Tensor) – \(k_i \in \{k_x,k_y,k_z\}\)

  • f (torch.Tensor) – A scalar function

Returns:

Gradient component

Return type:

torch.Tensor

Gradient Squared#

functional_tools.grad_dot_grad(kx, ky, kz, f)#

Computes squared gradient

This is a utility function that computes the squared gradient or the dot product of a gradient with itself

\[|\nabla f|^2 = \left(\frac{\partial f}{\partial x}\right)^2 + \left(\frac{\partial f}{\partial y}\right)^2 + \left(\frac{\partial f}{\partial z}\right)^2\]

where \(f\) is a given function.

Parameters:

  • ki (torch.Tensor) – \(k_i \in \{k_x,k_y,k_z\}\)

  • f (torch.Tensor) – A scalar function

Returns:

Dot product of gradient with itself

Return type:

torch.Tensor

Laplacian#

functional_tools.laplacian(k2, f)#

Computes Laplacian

This is a utility function that computes the Laplacian

\[\nabla^2 f = \frac{\partial^2 f}{\partial^2 x} + \frac{\partial^2 f}{\partial^2 y} + \frac{\partial^2 f}{\partial^2 z}\]

where \(f\) is a given function.

Parameters:

  • k2 (torch.Tensor) – \(k^2 = k_x^2 + k_y^2 + k_z^2\)

  • f (torch.Tensor) – A scalar function

Returns:

Laplacian

Return type:

torch.Tensor

Density Descriptors#

Reduced Gradient#

functional_tools.reduced_gradient(kx, ky, kz, den)#

Computes reduced gradient

This is a utility function that computes the reduced gradient

\[s = \frac{|\nabla n|}{2(3\pi)^{1/3} n^{4/3}}\]

where \(n\) is the electron density.

Parameters:

  • ki (torch.Tensor) – \(k_i \in \{k_x,k_y,k_z\}\)

  • den (torch.Tensor) – Electron density, \(n\)

Returns:

Reduced gradient

Return type:

torch.Tensor

Reduced Gradient Squared#

functional_tools.reduced_gradient_squared(kx, ky, kz, den)#

Computes squared reduced gradient

This is a utility function that computes the reduced gradient

\[s^2 = \frac{|\nabla n|^2}{4(3\pi)^{2/3} n^{8/3}}\]

where \(n\) is the electron density.

Parameters:

  • ki (torch.Tensor) – \(k_i \in \{k_x,k_y,k_z\}\)

  • den (torch.Tensor) – Electron density, \(n\)

Returns:

Squared reduced gradient

Return type:

torch.Tensor

Reduced Laplacian#

functional_tools.reduced_laplacian(k2, den)#

Computes reduced Laplacian

This is a utility function that computes the reduced Laplacian

\[q = \frac{\nabla^2 n}{4(3\pi)^{2/3} n^{5/3}}\]

where \(n\) is the electron density.

Parameters:

  • k2 (torch.Tensor) – \(k^2 = k_x^2 + k_y^2 + k_z^2\)

  • den (torch.Tensor) – Electron density, \(n\)

Returns:

Reduced Laplacian

Return type:

torch.Tensor

On this page
Show Source