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

Examples using torchpme.lib.KSpaceFilter

Examples of the KSpaceFilter class

Examples of the KSpaceFilter class

Custom models with automatic differentiation

Custom models with automatic differentiation

Splined potentials

Splined potentials

Computing LODE descriptors

Computing LODE descriptors