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

`__init__`	Initialize wavefunctions of specified size or with given coefficients.
`band_ke`	Return per-band kinetic energy of wavefunctions.
`band_norm`	Return per-band norm of wavefunctions.
`band_spin`	Return per-band spin of wavefunctions (must be spinorial).
`cat`	Join wavefunctions along specified dimension (default: 2 => bands).
`clone`	Create an independent copy (not a view / reference).
`constrain`	Enforce basis constraints on wavefunction coefficients.
`dot`	Inner (dot) product of self with other.
`dot_O`	Dot product of self with O(other).
`matmul`	Compute a matrix transformation (such as rotation) along the band dimension of wavefunction (self).
`n_bands`	Get number of bands in wavefunction
`norm`	Return overall norm of wavefunctions
`orthonormalize`	Return orthonormalized version of wavefunctions.
`overlap`	Return wavefunction with overlap operator applied.
`randomize`	Set wavefunction coefficients to bandwidth-limited random numbers.
`randomize_selected`	Randomize wavefunction coefficients of selected bands.
`read`	Read wavefunctions from cp_path.
`requires_grad_`	Set whether gradient with respect to this object is needed.
`split_bands`	Return wavefunction split by bands, bringing all basis coefficients of each band together on some process.
`split_basis`	Return wavefunction split by basis, bringing all bands of each basis coefficient together on some process.
`write`	Write wavefunctions to cp_path.
`zeros_like`	Create a zero Wavefunction similar to the present one.

Attributes

`non_spinor`	Return a non-spinorial view of this wavefunction.
`proj`	Projection of wavefunction on current projector in `Ions`.
`requires_grad`	Return whether gradient with respect to this object is needed.
`spinorial`	Whether this wavefunction has spinor components.
`basis`	Corresponding basis
`coeff`	Wavefunction coefficients (n_spins x nk x n_bands x n_spinor x n_basis)
`band_division`	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:

Wavefunction

clone()

Create an independent copy (not a view / reference).

Return type:: Wavefunction

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:: Wavefunction

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:: Wavefunction

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:: Waitable[Wavefunction]

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:: Waitable[Wavefunction]

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:: Wavefunction

band_division: TaskDivision | None: If present, wavefunctions are split along bands instead of along the ‘home position’ of split along basis.

basis: Basis: Corresponding 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.