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
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 x in place using op over self.comm.
Apply gradient operator to wavefunction C
Apply kinetic energy (KE) operator to wavefunction C
Apply potential V to wavefunction C
Collect density contributions given wavefunction C and occupations f.
Number of FFTs to perform together.
Kinetic energy (KE) of each plane wave in basis in \(E_h\)
Kinetic energy (KE) of each plane wave in basis in \(E_h\)
Kinetic energy (KE) stress tensor of each plane wave in basis in \(E_h\)
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
Indexing datatype for pad_index and pad_index_mine
n_avg accounting for real-basis weights if any.
Ideal n_avg_weighted based on ke_cutoff G-sphere volume.
Basis/bands communicator
Overall k-points and basis/bands communicator
Lattice vectors of unit cell
implicit in basis for ultrasoft / PAW
Corresponding k-point set
Default number of spin channels
Default number of spinorial components
Subset of k handled by this basis (due to MPI division)
Subset of weights corresponding to k
Combined spin and k-point weights
Whether wavefunctions are real
Kinetic energy cutoff
Wavefunction grid (always process-local)
Plane waves in reciprocal lattice coordinates
Number of plane waves for each k
Minimum of n across all k (including on other processes)
Maximum of n across all k (including on other processes)
Average n across all k (weighted by wk)
Actual common n stored for each k including padding
Index of each plane wave in reciprocal grid
Number of bands to FFT together
Number of bands to MPI transfer together
Which basis entries are padding (beyond n)
Subset of pad_index on this process
Division of basis across rc.comm_b
Slice of basis entries local to this process
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:
- apply_ke(C)
Apply kinetic energy (KE) operator to wavefunction C
- Parameters:
self (Basis) –
C (Wavefunction) –
- Return type:
- apply_potential(V, C)
Apply potential V to wavefunction C
- Parameters:
self (Basis) –
V (FieldH) –
C (Wavefunction) –
- Return type:
- 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:
- 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
- iG: torch.Tensor
Plane waves in reciprocal lattice coordinates
- k: torch.Tensor
Subset of k handled by this basis (due to MPI division)
- ke_cutoff: float
Kinetic energy cutoff
- 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