qimpy.dft.electrons.Basis

class Basis(*, process_grid, lattice, ions, symmetries, kpoints, n_spins, n_spinor, checkpoint_in=(None, ''), ke_cutoff=20.0, real_wavefunctions=False, grid=None, fft_block_size=0, mpi_block_size=0)

Bases: TreeNode

Plane-wave basis for electronic wavefunctions. The underlying qimpy.mpi.TaskDivision splits plane waves over rc.comm_b

Parameters:

process_grid (ProcessGrid) –
lattice (Lattice) –
ions (Ions) –
symmetries (Symmetries) –
kpoints (Kpoints) –
n_spins (int) –
n_spinor (int) –
checkpoint_in (CheckpointPath) –
ke_cutoff (float) –
real_wavefunctions (bool) –
grid (Optional[Union[Grid, dict]]) –
fft_block_size (int) –
mpi_block_size (int) –

__init__(*, process_grid, lattice, ions, symmetries, kpoints, n_spins, n_spinor, checkpoint_in=(None, ''), ke_cutoff=20.0, real_wavefunctions=False, grid=None, fft_block_size=0, mpi_block_size=0)

Initialize plane-wave basis with ke_cutoff.

Parameters:

lattice (Lattice) – Lattice whose reciprocal lattice vectors define plane-wave basis.
ions (Ions) – Ions that specify the pseudopotential portion of the basis; the basis implicitly depends on the ion positions for ultrasoft or PAW due to the augmentation of all operators at each ion.
symmetries (Symmetries) – Symmetries with which the wavefunction grid should be commensurate.
kpoints (Kpoints) – Set of k-points to initialize basis for. Note that the basis is only initialized for k-points to be operated on by current process i.e. for k = kpoints.k[kpoints.i_start : kpoints.i_stop].
n_spins (int) – Default number of spin channels for wavefunctions in this basis.
n_spinor (int) – Default number of spinor components for wavefunctions in this basis. Also used only to check support for real_wavefunctions.
ke_cutoff (float) – [Input file] Wavefunction kinetic energy cutoff in Hartrees.
real_wavefunctions (bool) – [Input file] Whether to use real wavefunctions (instead of complex). This is only supported for non-spinorial, Gamma-point-only calculations, where conjugate symmetry allows real wavefunctions.
grid (Grid | dict | None) – [Input file] Override parameters of grid for wavefunction operations.
fft_block_size (int) – [Input file] Number of wavefunction bands to FFT simultaneously. Higher numbers require more memory, but can achieve better occupancy of GPUs or high-core-count CPUs. The default of 0 auto-selects the block size based on the number of bands and k-points being processed by each process.
mpi_block_size (int) – [Input file] Number of wavefunction bands to MPI transfer simultaneously. Lower numbers may allow better overlap between computation and transfers, which is beneficial if MPI implementation supports asynchronous progress and/or CUDA streams are used to compute asynrchronously. Higher numbers mitigate MPI latency, but may require more memory. This number is automatically rounded up to nearest multiple of fft_block_size * comm.size. The default of 0 selects the block size based on the number of bands and k-points being processed by each process.
process_grid (ProcessGrid) –
checkpoint_in (CheckpointPath) –

Return type:

None

Methods

`__init__`	Initialize plane-wave basis with ke_cutoff.
`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.
`allreduce_in_place`	Allreduce x in place using op over self.comm.
`apply_gradient`	Apply gradient operator to wavefunction C
`apply_ke`	Apply kinetic energy (KE) operator to wavefunction C
`apply_potential`	Apply potential V to wavefunction C
`collect_density`	Collect density contributions given wavefunction C and occupations f.
`get_fft_block_size`	Number of FFTs to perform together.
`get_gradient`	Kinetic energy (KE) of each plane wave in basis in \(E_h\)
`get_ke`	Kinetic energy (KE) of each plane wave in basis in \(E_h\)
`get_ke_stress`	Kinetic energy (KE) stress tensor of each plane wave in basis in \(E_h\)
`get_mpi_block_size`	Number of bands to MPI transfer together for collect_density and apply_potential.
`save_checkpoint`	Save self and all children in hierarchy to cp_path.

Attributes

`PadIndex`	Indexing datatype for pad_index and pad_index_mine
`n_avg_weighted`	n_avg accounting for real-basis weights if any.
`n_ideal`	Ideal n_avg_weighted based on ke_cutoff G-sphere volume.
`comm`	Basis/bands communicator
`comm_kb`	Overall k-points and basis/bands communicator
`lattice`	Lattice vectors of unit cell
`ions`	implicit in basis for ultrasoft / PAW
`kpoints`	Corresponding k-point set
`n_spins`	Default number of spin channels
`n_spinor`	Default number of spinorial components
`k`	Subset of k handled by this basis (due to MPI division)
`wk`	Subset of weights corresponding to k
`w_sk`	Combined spin and k-point weights
`real_wavefunctions`	Whether wavefunctions are real
`ke_cutoff`	Kinetic energy cutoff
`grid`	Wavefunction grid (always process-local)
`iG`	Plane waves in reciprocal lattice coordinates
`n`	Number of plane waves for each k
`n_min`	Minimum of n across all k (including on other processes)
`n_max`	Maximum of n across all k (including on other processes)
`n_avg`	Average n across all k (weighted by wk)
`n_tot`	Actual common n stored for each k including padding
`fft_index`	Index of each plane wave in reciprocal grid
`fft_block_size`	Number of bands to FFT together
`mpi_block_size`	Number of bands to MPI transfer together
`pad_index`	Which basis entries are padding (beyond n)
`pad_index_mine`	Subset of pad_index on this process
`division`	Division of basis across rc.comm_b
`mine`	Slice of basis entries local to this process
`real`	Extra indices for real wavefunctions
`child_names`	Names of attributes with child objects.

allreduce_in_place(x, op=<mpi4py.MPI.Op object>)

Allreduce x in place using op over self.comm. Convenient wrapper used in many basis operations.

Parameters:

x (Tensor) –
op (Op) –

Return type:

None

apply_gradient(C, i_dir)

Apply gradient operator to wavefunction C

Parameters:

self (Basis) –
C (Wavefunction) –
i_dir (int) –

Return type:

Wavefunction

apply_ke(C)

Apply kinetic energy (KE) operator to wavefunction C

Parameters:

self (Basis) –
C (Wavefunction) –

Return type:

Wavefunction

apply_potential(V, C)

Apply potential V to wavefunction C

Parameters:

self (Basis) –
V (FieldH) –
C (Wavefunction) –

Return type:

Wavefunction

collect_density(C, f, need_Mvec)

Collect density contributions given wavefunction C and occupations f. The result is in real-space on basis.grid. If the wavefunction has two spin channels, the two components of the resulting density correspond to total density and magentization. If need_Mvec = True i.e. magnetization vector (C must be spinorial for this), the result contains 4 components: (n_tot, Mx, My, Mz). Here, (n_tot +/- Mz)/2 yield the \(\rho_{\uparrow\uparrow}\) and \(\rho_{\downarrow\downarrow}\) components of the spin density matrix, while (Mx +/- i My)/2 yield the \(\rho_{\uparrow\downarrow}\) and \(\rho_{\downarrow\uparrow}\) components of the spin density matrix.

Parameters:

self (Basis) –
C (Wavefunction) –
f (Tensor) –
need_Mvec (bool) –

Return type:

FieldR

get_fft_block_size(n_batch, n_bands)

Number of FFTs to perform together. Equals fft_block_size, if that is non-zero, and uses a heuristic based on batch dimension and number of bands.

Parameters:

n_batch (int) –
n_bands (int) –

Return type:

int

get_gradient(basis_slice=slice(None, None, None))

Kinetic energy (KE) of each plane wave in basis in \(E_h\)

Parameters:: basis_slice (slice) – Selection of basis functions to get KE for (default: full basis)
Returns:: TODO, dimensions: nk_mine x len(basis_slice)
Return type:: torch.Tensor

get_ke(basis_slice=slice(None, None, None))

Kinetic energy (KE) of each plane wave in basis in \(E_h\)

Parameters:: basis_slice (slice) – Selection of basis functions to get KE for (default: full basis)
Returns:: KE for each plane-wave, dimensions: nk_mine x len(basis_slice)
Return type:: torch.Tensor

get_ke_stress(basis_slice=slice(None, None, None))

Kinetic energy (KE) stress tensor of each plane wave in basis in \(E_h\)

Parameters:

basis_slice (slice) – Selection of basis functions to get KE for (default: full basis)

Returns:

KE tensor for each plane-wave,: dimensions: nk_mine x len(basis_slice) x 3 x 3

Return type:

torch.Tensor

get_mpi_block_size(n_batch, n_bands, fft_block_size)

Number of bands to MPI transfer together for collect_density and apply_potential. Uses mpi_block_size, if that is non-zero, and uses a heuristic based on batch dimension and number of bands. The final number is coerced to a multiple of fft_block_size * comm.size or rounded up to n_bands, if it is already close to that limit.

Parameters:

n_batch (int) –
n_bands (int) –
fft_block_size (int) –

Return type:

int

PadIndex

Indexing datatype for pad_index and pad_index_mine

alias of tuple[slice, Tensor, slice, slice, Tensor]

comm: MPI.Comm: Basis/bands communicator

comm_kb: MPI.Comm: Overall k-points and basis/bands communicator

division: TaskDivision: Division of basis across rc.comm_b

fft_block_size: int: Number of bands to FFT together

fft_index: torch.Tensor: Index of each plane wave in reciprocal grid

grid: Grid: Wavefunction grid (always process-local)

iG: torch.Tensor: Plane waves in reciprocal lattice coordinates

ions: Ions

implicit in basis for ultrasoft / PAW

Type:: Ionic system

k: torch.Tensor: Subset of k handled by this basis (due to MPI division)

ke_cutoff: float: Kinetic energy cutoff

kpoints: Kpoints: Corresponding k-point set

lattice: Lattice: Lattice vectors of unit cell

mine: slice: Slice of basis entries local to this process

mpi_block_size: int: Number of bands to MPI transfer together

n: torch.Tensor: Number of plane waves for each k

n_avg: float: Average n across all k (weighted by wk)

property n_avg_weighted: float: n_avg accounting for real-basis weights if any.

property n_ideal: float: Ideal n_avg_weighted based on ke_cutoff G-sphere volume.

n_max: int: Maximum of n across all k (including on other processes)

n_min: int: Minimum of n across all k (including on other processes)

n_spinor: int: Default number of spinorial components

n_spins: int: Default number of spin channels

n_tot: int: Actual common n stored for each k including padding

pad_index: PadIndex: Which basis entries are padding (beyond n)

pad_index_mine: PadIndex: Subset of pad_index on this process

real: BasisReal: Extra indices for real wavefunctions

real_wavefunctions: bool: Whether wavefunctions are real

w_sk: torch.Tensor: Combined spin and k-point weights

wk: torch.Tensor: Subset of weights corresponding to k