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 (TaskDivision | None) –
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 (Tensor | None) – 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 (TaskDivision | None) – 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.
grad
optional gradient (of energy) with respect to this object.
- 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 (int | None) – 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: TaskDivision | None
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.