qimpy.grid.Grid

class Grid(*, lattice, symmetries, group, allow_parallel=True, checkpoint_in=(None, ''), ke_cutoff_wavefunction=None, ke_cutoff=None, shape=None)

Bases: TreeNode

Real and reciprocal space grids for a unit cell. The grid could either be local or distributed over a process group, and this class provides FFT routines to switch fields on these grids, and routines to convert fields between grids.

Parameters:

lattice (Lattice)
symmetries (Symmetries)
group (ProcessGroup | None)
allow_parallel (bool)
checkpoint_in (CheckpointPath)
ke_cutoff_wavefunction (float | None)
ke_cutoff (float)
shape (tuple[int, ...])

__init__(*, lattice, symmetries, group, allow_parallel=True, checkpoint_in=(None, ''), ke_cutoff_wavefunction=None, ke_cutoff=None, shape=None)

Create local or distributed grid for lattice.

Parameters:

lattice (Lattice) – Lattice whose reciprocal lattice vectors define plane-wave basis
symmetries (Symmetries) – Symmetries with which grid dimensions will be made commensurate, checked if specified explicitly by shape below and used for symmetrization of Field’s associated with this grid.
group (ProcessGroup | None) – Process group to split grid (and its FFTs) over, if provided. This is only used if allow_parallel is set to True.
allow_parallel (bool) – [Input file] Whether to allow splitting the grid over a process group. If set to False, the grid is not split regardless of whether a process group is provided. This may be useful to avoid overhead from parallel FFTs in large GPU jobs.
ke_cutoff_wavefunction (float | None) – Plane-wave kinetic-energy cutoff in \(E_h\) for any electronic wavefunctions to be used with this grid. This is an internally set parameter that effectively sets the default for ke_cutoff to 4 * ke_cutoff_wavefunction, and has no effect if either ke_cutoff or shape is specified explicitly.
ke_cutoff (float | None) – [Input file] Kinetic-energy cutoff for grid in Hartrees. If unspecified, this is taken to be its minimum value: 4x the wavefunction kinetic energy cutoff. This has no effect if shape is specified explicitly.
shape (Sequence[int | None]) – [Input file] Explicit grid dimensions [Nx, Ny, Nz]. Highest precedence, and if specified, will supercede ke_cutoff.
checkpoint_in (CheckpointPath)

Return type:

None

Methods

`__init__`	Create local or distributed grid for lattice.
`add_child`	Construct child object self.`attr_name` of type cls.
`add_child_one_of`	Invoke add_child on one of several child options in args.
`fft`	Forward Fast Fourier Transform.
`get_Gmax`	Get maximum wave-vector magnitude of the FFT grid.
`get_gradient_operator`	Get gradient operator in reciprocal space.
`get_mesh`	Get mesh integer coordinates for real or reciprocal space
`get_nyquist_indices`	List of dimensions and corresponding indices of Nyquist frequencies.
`ifft`	Inverse Fast Fourier Transform.
`is_equivalent_to`	Whether this grid is equivalent to grid.
`save_checkpoint`	Save self and all children in hierarchy to cp_path.
`zero_nyquist`	Set Nyquist components of x corresponding to space 'G' or 'H' to zero.

Attributes

`dV`	Volume per grid point (real-space integration weight).
`field_symmetrizer`	Symmetrizer for fields on this grid (initialized on first use).
`weight2H`	Hermitian-symmetry weights for reduced reciprocal space.
`lattice`
`symmetries`
`group`	Process group to split grid and FFTs over
`n_procs`	Size of proces group
`i_proc`	Rank within process group
`is_split`	Whether the grid is split over MPI
`ke_cutoff`	Kinetic energy of Nyquist-frequency plane-waves
`shape`	Global real-space grid dimensions
`shapeH`	Global half-reciprocal-space grid dimensions
`shapeR_mine`	Local real grid dimensions
`shapeG_mine`	Local reciprocal grid dimensions
`shapeH_mine`	Local half-reciprocal grid dimensions
`split0`	MPI split of real-space dimension 0
`split2`	MPI split of reciprocal dimension 2
`split2H`	MPI split of half-reciprocal dimension 2
`child_names`	Names of attributes with child objects.
`variant_name`	Version of children having variants (if any)

fft(v)

Forward Fast Fourier Transform. This method dispatches to complex-to-complex or real-to-complex transforms depending on whether the input v is complex or real. Note that QimPy applies normalization in forward transforms, corresponding to norm=’forward’ in the torch.fft routines. This makes the G=0 components in reciprocal space correspond to the mean value of the real space version.

Parameters:: v (torch.Tensor (complex or real)) – Last 3 dimensions must match shapeR_mine, and any preceding dimensions are batched over.
Returns:: Last 3 dimensions will be shapeG_mine or shapeH_mine, depending on whether v is complex or real respectively, preceded by any batch dimensions in the input.
Return type:: torch.Tensor (complex)

get_Gmax()

Get maximum wave-vector magnitude of the FFT grid.

Return type:: float

get_gradient_operator(space, zero_nyquist)

Get gradient operator in reciprocal space.

Parameters:

space ({'G', 'H'}) – Which space to compute mesh coordinates for: ‘G’ = full reciprocal space and ‘H’ = half or Hermitian-symmetric recipocal space.
zero_nyquist (bool) – Whether to zero out Nyquist frequency components

Returns:

Tensor with dimensions (3,) + shape_mine, where shape_mine is the relevant local dimensions of requested space

Return type:

Tensor

get_mesh(space, mine=True)

Get mesh integer coordinates for real or reciprocal space

Parameters:

space ({'R', 'G', 'H'}) – Which space to compute mesh coordinates for: ‘R’ = real space, ‘G’ = full reciprocal space and ‘H’ = half or Hermitian-symmetric recipocal space resulting from FFT of real functions on grid.
mine (bool) – Only get local portion of mesh if True (default). Otherwise, get the entire mesh, even if grid is split over MPI.

Returns:

Integer tensor with dimensions shape_mine + (3,), where shape_mine is the relevant local dimensions of requested space if mine is True (and with shape, instead if mine is False).

Return type:

Tensor

get_nyquist_indices(space)

List of dimensions and corresponding indices of Nyquist frequencies.

Parameters:: space (str)
Return type:: list[tuple(int, int)]

ifft(v)

Inverse Fast Fourier Transform. This method dispatches to complex-to-complex or complex-to-real transforms depending on whether the last three dimensions of v match shapeG_mine or shapeH_mine respectively. Note that QimPy applies normalization in forward transforms (see qimpy.grid.Grid.fft()).

Parameters:: v (torch.Tensor (complex)) – Last 3 dimensions must match shapeG_mine, and any preceding dimensions are batched over
Returns:: Last 3 dimensions will be shapeR_mine, preceded by any batch dimensions in the input. The result will be complex or real, depending on whether the last three dimensions of v match shapeG_mine or shapeH_mine.
Return type:: torch.Tensor (complex or real)

is_equivalent_to(grid)

Whether this grid is equivalent to grid.

Parameters:: grid (Grid)
Return type:: bool

zero_nyquist(x, space)

Set Nyquist components of x corresponding to space ‘G’ or ‘H’ to zero.

Parameters:

x (Tensor)
space (str)

Return type:

None

property dV: float: Volume per grid point (real-space integration weight).

property field_symmetrizer: FieldSymmetrizer: Symmetrizer for fields on this grid (initialized on first use).

group: ProcessGroup | None: Process group to split grid and FFTs over

i_proc: int: Rank within process group

is_split: bool: Whether the grid is split over MPI

ke_cutoff: float: Kinetic energy of Nyquist-frequency plane-waves

n_procs: int: Size of proces group

shape: tuple[int, ...]: Global real-space grid dimensions

shapeG_mine: tuple[int, ...]: Local reciprocal grid dimensions

shapeH: tuple[int, ...]: Global half-reciprocal-space grid dimensions

shapeH_mine: tuple[int, ...]: Local half-reciprocal grid dimensions

shapeR_mine: tuple[int, ...]: Local real grid dimensions

split0: TaskDivision: MPI split of real-space dimension 0

split2: TaskDivision: MPI split of reciprocal dimension 2

split2H: TaskDivision: MPI split of half-reciprocal dimension 2

property weight2H: Tensor: Hermitian-symmetry weights for reduced reciprocal space.