Kspace filter

class torchpme.lib.KSpaceKernel[source]

Base class defining the interface for a reciprocal-space kernel helper.

Provides an interface to compute the reciprocal-space convolution kernel that is used e.g. to compute potentials using Fourier transforms. Parameters of the kernel in derived classes should be defined and stored in the __init__ method.

NB: we need this slightly convoluted way of implementing what often amounts to a simple, pure function of \(|\mathbf{k}|^2\) in order to be able to provide a customizable filter class that can be jitted.

kernel_from_k_sq(kvectors: Tensor) → Tensor[source]

Computes the reciprocal-space kernel on a grid of k points given a tensor containing \(\mathbf{k}\).

Parameters:

kvectors (Tensor) – torch.tensor containing the k vector at which the kernel is to be evaluated.

Return type:

Tensor

class torchpme.lib.KSpaceFilter(cell: Tensor, ns_mesh: Tensor, kernel: KSpaceKernel, fft_norm: str = 'ortho', ifft_norm: str = 'ortho')[source]

Apply a reciprocal-space filter to a real-space mesh.

The class combines the costruction of a reciprocal-space grid \(\{\mathbf{k}_n\}\) (that should be commensurate to the grid in real space, so the class takes the same options as MeshInterpolator), the calculation of a scalar filter function \(\phi(|\mathbf{k}|^2)\), defined as a function of the squared norm of the reciprocal space grid points, and the application of the filter to a real-space function \(f(\mathbf{x})\), defined on a mesh \(\{\mathbf{x}_n\}\).

In practice, the application of the filter amounts to \(f\rightarrow \hat{f} \rightarrow \hat{\tilde{f}}= \hat{f} \phi \rightarrow \tilde{f}\)

See also the Examples of the KSpaceFilter class for a demonstration of the functionalities of this class.

Parameters:

  • cell (Tensor) – torch.tensor of shape (3, 3), where cell[i] is the i-th basis vector of the unit cell

  • ns_mesh (Tensor) – toch.tensor of shape (3,) Number of mesh points to use along each of the three axes

  • kernel (KSpaceKernel) – KSpaceKernel A KSpaceKernel-derived class providing a from_k_sq method that evaluates \(\psi\) given the square modulus of the k-space mesh points

  • fft_norm (str) – str The normalization applied to the forward FT. Can be “forward”, “backward”, “ortho”. See torch:fft:rfftn()

  • ifft_norm (str) – str The normalization applied to the inverse FT. Can be “forward”, “backward”, “ortho”. See torch:fft:irfftn()

update(cell: Tensor | None = None, ns_mesh: Tensor | None = None) → None[source]

Update buffers and derived attributes of the instance.

If neither cell nor ns_mesh are passed, only the filter is updated, typically following a change in the underlying potential. If cell and/or ns_mesh are given, the instance’s attributes required by these will also be updated accordingly.

Parameters:

  • cell (Tensor | None) – torch.tensor of shape (3, 3), where cell[i] is the i-th basis vector of the unit cell

  • ns_mesh (Tensor | None) – toch.tensor of shape (3,) Number of mesh points to use along each of the three axes

Return type:

None

forward(mesh_values: Tensor) → Tensor[source]

Applies the k-space filter by Fourier transforming the given mesh_values tensor, multiplying the result by the filter array (that should have been previously computed with a call to update()) and Fourier-transforming back to real space.

If you update the cell, the ns_mesh or anything inside the kernel object after you initlized the object, you have call update() to update the object calling this method.

kernel_filter.update(cell)
kernel_filter.forward(mesh)
Parameters:

mesh_values (Tensor) – torch.tensor of shape (n_channels, nx, ny, nz) The values of the input function on a real-space mesh. Shape should match the shape of the filter.

Returns:

torch.tensor of shape (n_channels, nx, ny, nz) The real-space mesh containing the transformed function values.

Return type:

Tensor

class torchpme.lib.P3MKSpaceFilter(cell: Tensor, ns_mesh: Tensor, interpolation_nodes: int, kernel: KSpaceKernel, fft_norm: str = 'ortho', ifft_norm: str = 'ortho', mode: int = 0, differential_order: int = 2)[source]

A specialized implementation of the k-space filter for the P3M method, with a cell-dependent Green’s function kernel. This class does almost the same thing as KSpaceFilter, but with a different, P3M-specialized filter. See this paper for your reference.

Parameters:

  • cell (Tensor) – torch.tensor of shape (3, 3), where cell[i] is the i-th basis vector of the unit cell

  • ns_mesh (Tensor) – toch.tensor of shape (3,) Number of mesh points to use along each of the three axes

  • interpolation_nodes (int) – int The number n of nodes used in the interpolation per coordinate axis. The total number of interpolation nodes in 3D will be n^3. In general, for n nodes, the interpolation will be performed by piecewise polynomials of degree n - 1 (e.g. n = 4 for cubic interpolation). Only the values 1, 2, 3, 4, 5 are supported.

  • kernel (KSpaceKernel) – KSpaceKernel A KSpaceKernel-derived class providing a from_k_sq method that evaluates \(\psi\) given the square modulus of the k-space mesh points

  • fft_norm (str) – str The normalization applied to the forward FT. Can be “forward”, “backward”, “ortho”. See torch:fft:rfftn()

  • ifft_norm (str) – str The normalization applied to the inverse FT. Can be “forward”, “backward”, “ortho”. See torch:fft:irfftn()

  • mode (int) – int, 0 for the electrostatic potential, 1 for the electrostatic energy, 2 for the dipolar torques, and 3 for the dipolar forces. For more details, see eq.30 of that paper.

  • diff_order –

    int, the order of the approximation of the difference operator. Higher order is more accurate, but also more expensive. For more details, see Appendix C of this paper. The values 1, 2, 3, 4, 5, 6 are supported.

  • differential_order (int)

update(cell: Tensor | None = None, ns_mesh: Tensor | None = None) → None[source]

Update buffers and derived attributes of the instance.

If neither cell nor ns_mesh are passed, only the filter is updated, typically following a change in the underlying potential. If cell and/or ns_mesh are given, the instance’s attributes required by these will also be updated accordingly.

Parameters:

  • cell (Tensor | None) – torch.tensor of shape (3, 3), where cell[i] is the i-th basis vector of the unit cell

  • ns_mesh (Tensor | None) – toch.tensor of shape (3,) Number of mesh points to use along each of the three axes

Return type:

None