qimpy.dft.electrons.Wavefunction
- class Wavefunction(basis, *, coeff=None, band_division=None, n_bands=0, n_spins=0, n_spinor=0)
Bases:
Gradable
[Wavefunction
]Electronic wavefunctions including coefficients and projections
- Parameters
basis (Basis) –
coeff (Tensor) –
band_division (Optional[TaskDivision]) –
n_bands (int) –
n_spins (int) –
n_spinor (int) –
- __init__(basis, *, coeff=None, band_division=None, n_bands=0, n_spins=0, n_spinor=0)
Initialize wavefunctions of specified size or with given coefficients.
- Parameters
basis (Basis) – Basis set for wavefunction
coeff (Optional[Tensor]) – Wavefunction coefficients in specified basis. If not provided, one of band_division or n_bands must be provided to initialize zero wavefunctions of determined size.
band_division (Optional[TaskDivision]) – If None, wavefunctions are MPI-split over basis coefficients, which is the default or “home” position for the wavefunctions. When not None, the wavefunctions are split over bands as specified by the TaskDivision object, with all basis coefficients for each band together on some process.
n_bands (int, default: band_division.n_mine if available, else 0) – Number of bands to initialize zero wavefunctions for. Used only if coeff is None, and this does not initialize proj.
n_spins (int) – If non-zero, override the number of spin channels in new wavefunctions (instead of the value in basis). Used only if coeff is None, and n_bands or band_division is given
n_spinor (int) – If non-zero, override the number of spinor components in new wavefunctions (instead of the value in basis). Used only if coeff is None, and n_bands or band_division is given
- Return type
None
Methods
Initialize wavefunctions of specified size or with given coefficients.
Return per-band kinetic energy of wavefunctions.
Return per-band norm of wavefunctions.
Return per-band spin of wavefunctions (must be spinorial).
Join wavefunctions along specified dimension (default: 2 => bands).
Create an independent copy (not a view / reference).
Enforce basis constraints on wavefunction coefficients.
Inner (dot) product of self with other.
Dot product of self with O(other).
Compute a matrix transformation (such as rotation) along the band dimension of wavefunction (self).
Get number of bands in wavefunction
Return overall norm of wavefunctions
Return orthonormalized version of wavefunctions.
Return wavefunction with overlap operator applied.
Set wavefunction coefficients to bandwidth-limited random numbers.
Randomize wavefunction coefficients of selected bands.
Read wavefunctions from cp_path.
requires_grad_
Set whether gradient with respect to this object is needed.
Return wavefunction split by bands, bringing all basis coefficients of each band together on some process.
Return wavefunction split by basis, bringing all bands of each basis coefficient together on some process.
Write wavefunctions to cp_path.
Create a zero Wavefunction similar to the present one.
Attributes
Return a non-spinorial view of this wavefunction.
Projection of wavefunction on current projector in
Ions
.requires_grad
Return whether gradient with respect to this object is needed.
Whether this wavefunction has spinor components.
Corresponding basis
Wavefunction coefficients (n_spins x nk x n_bands x n_spinor x n_basis)
If present, wavefunctions are split along bands instead of along the 'home position' of split along basis.
- band_ke()
Return per-band kinetic energy of wavefunctions.
- Parameters
self (Wavefunction) –
- Return type
Tensor
- band_norm()
Return per-band norm of wavefunctions.
- Parameters
self (Wavefunction) –
- Return type
Tensor
- band_spin()
Return per-band spin of wavefunctions (must be spinorial). Result dimensions are 3 x 1 x nk x n_bands.
- Parameters
self (Wavefunction) –
- Return type
Tensor
- cat(other, dim=2, clear=False)
Join wavefunctions along specified dimension (default: 2 => bands). If clear is True, eagerly clear memory of the input operands. Note that this will leave self and other in a broken state, so use this only if they are deleted or replaced shortly thereafter. Despite this danger, this is often necessary becaause this operation will likely be near the peak memory usage eg. within Davidson and CheFSI.
- Parameters
self (Wavefunction) –
other (Wavefunction) –
dim (int) –
clear (bool) –
- Return type
- clone()
Create an independent copy (not a view / reference).
- Return type
- constrain()
Enforce basis constraints on wavefunction coefficients. This includes setting padded coefficients to zero, and imposing Hermitian symmetry in Gz = 0 coefficients for real wavefunctions.
- Parameters
self (Wavefunction) –
- Return type
None
- dot(other)
Inner (dot) product of self with other. For convenience, this can also be invoked as self ^ other. Note that this means xor(self, other) also computes wavefunction dot instead of logical xor (which is meaningless for wavefunctions anyway). All dimensions except bands must match between self and other.
- Parameters
other (qimpy.electrons.Wavefunction) – Dimensions must match self for basis, can differ for bands, and must be broadcastable for preceding dimensions (spin and k). If the spinor dimensions are different, as utilized for non-spinorial projectors in spinorial calculations, the spinor dimenion is not contracted over and instead integrated as adjacent bands of the non-spinorial input.
overlap (bool, default: False) – If True, include the overlap operator in the product. The overlap operator is identity for norm-conserving pseudopotentials, but includes augmentation for ultrasoft and PAW pseudopotentials
self (Wavefunction) –
- Returns
Dimensions: n_spins x nk_mine x n_bands(self) x n_bands(other). If only one of the two inputs is spinorial, then n_bands of the non-spinorial one is effectively doubled, including spinor components in adjacent entries.
- Return type
torch.Tensor
- dot_O(other)
Dot product of self with O(other). Here, the overlap operator O is identity for norm-conserving pseudopotentials, but includes augmentation for ultrasoft and PAW pseudopotentials.
- Returns
Dimensions: n_spins x nk_mine x n_bands(self) x n_bands(other)
- Return type
torch.Tensor
- Parameters
self (Wavefunction) –
other (Wavefunction) –
- matmul(mat)
Compute a matrix transformation (such as rotation) along the band dimension of wavefunction (self). For convenience, this can also be invoked as self @ other, using the standard matrix multiply operator.
- Parameters
mat (torch.Tensor) – Last two dimensions specify transformation in band space: penultimate dimension should match n_bands of the wavefunction and final dimension determines n_bands of output. Preceding dimensions should be broadcastable with spin and k. For the special case of non-spinorial projectors in spinorial calculations, penultimate dimension will be twice of n_bands.
self (Wavefunction) –
- Returns
The result will have n_bands = mat.shape[-1], same basis as input, and spin and k determined by broadcasting with mat.shape[:-2]. The output will have the same spinor components as input, except for non-spinorial projectors in spinorial calculations, where the result will be spinorial.
- Return type
qimpy.electrons.Wavefunction
- n_bands()
Get number of bands in wavefunction
- Return type
int
- norm()
Return overall norm of wavefunctions
- Parameters
self (Wavefunction) –
- Return type
float
- orthonormalize()
Return orthonormalized version of wavefunctions. This internally uses Gram-Schmidt scheme using a Cholesky decomposition, which is not differentiable. Use a symmetric orthonormalization scheme using
qimpy.mpi.ortho_matrix()
for a differentiable scheme.- Parameters
self (Wavefunction) –
- Return type
- overlap()
Return wavefunction with overlap operator applied. This is identity and returns a view of the original wavefunction for norm-conserving pseudopotentials, but includes augmentation terms for ultrasoft and PAW pseudopotentials
- Parameters
self (Wavefunction) –
- Return type
- randomize(seed=0, b_start=0, b_stop=None)
Set wavefunction coefficients to bandwidth-limited random numbers. This is done reproducibly, regardless of MPI configuration of the run, by running a separate xor-shift random number generator with a different seed at each combination of spin, k, spinor and G-vector, looping only over the bands to generate the random number.
- Parameters
seed (int) – Seed offset at each k and iG for xor-shift random number generator
b_start (int) – Starting band index (global index if MPI-split over bands)
b_stop (Optional[int]) – Stopping band index (global index if MPI-split over bands)
self (Wavefunction) –
- Return type
None
- randomize_selected(i_spin, i_k, i_band, seed)
Randomize wavefunction coefficients of selected bands. The bands are indexed by a tuple (i_spin, i_k, i_band) of 1D index tensors with same shape. This is only supported for wavefunctions split over basis (i.e. no band_division).
- Parameters
self (Wavefunction) –
i_spin (Tensor) –
i_k (Tensor) –
i_band (Tensor) –
seed (int) –
- Return type
None
- read(cp_path)
Read wavefunctions from cp_path. Return number of bands read.
- Parameters
cp_path (CheckpointPath) –
- Return type
int
- split_bands()
Return wavefunction split by bands, bringing all basis coefficients of each band together on some process. Note that the result may be a view if there is only one process, or if the wavefunction is already split by bands
- Parameters
self (Wavefunction) –
- Return type
- split_basis()
Return wavefunction split by basis, bringing all bands of each basis coefficient together on some process. Note that the result may be a view if there is only one process, or if the wavefunction is already split by basis
- Parameters
self (Wavefunction) –
- Return type
- write(cp_path)
Write wavefunctions to cp_path.
- Parameters
cp_path (CheckpointPath) –
- Return type
None
- zeros_like(n_bands=0)
Create a zero Wavefunction similar to the present one. Optionally override the number of bands with n_bands.
- Parameters
n_bands (int) –
- Return type
- band_division: Optional[TaskDivision]
If present, wavefunctions are split along bands instead of along the ‘home position’ of split along basis.
- coeff: torch.Tensor
Wavefunction coefficients (n_spins x nk x n_bands x n_spinor x n_basis)
- property non_spinor: Wavefunction
Return a non-spinorial view of this wavefunction. If spinorial, the result will have twice the bands instead. This is useful when dealing with projectors which remain non-spinorial, even in spinorial calculations. The view is identical to the input for non-spinorial wavefunctions.
- property proj: Tensor
Projection of wavefunction on current projector in
Ions
. This is computed, cached and invalidated automatically.
- property spinorial: bool
Whether this wavefunction has spinor components.