Wavefunction / Basis

Wavefunction / Basis#

Basis#

This module contains everything related to the basis set. This includes the frequently used IndexHelper, which expedites transformations between atom-, shell-, and orbital-resolved properties.

Indexhelper#

Basis: IndexHelper#

Index helper utility to create index maps between atomic, shell-resolved, and orbital resolved representations of quantities.

Example

import torch
from dxtb import IndexHelper

# Define atomic numbers and angular momentum for each element
numbers = torch.tensor([6, 1, 1, 1, 1])
angular = {1: [0], 6: [0, 1]}

# Create an IndexHelper instance with angular momentum specifications
ihelp = IndexHelper.from_numbers_angular(numbers, angular)

# Count the number of entries in the angular momentum tensor
result = torch.sum(ihelp.angular >= 0)
print(result)  # torch.tensor(6)
class dxtb._src.basis.indexhelper.IndexHelper(unique_angular, angular, atom_to_unique, ushells_to_unique, ushells_per_unique, shells_to_ushell, shells_per_atom, shell_index, shells_to_atom, orbitals_per_shell, orbital_index, orbitals_to_shell, batch_mode, device=None, dtype=torch.int64, *, store=None, **_)[source]

Bases: TensorLike

Index helper for basis set.

Parameters:
cpu()

Returns a copy of the TensorLike instance on the CPU.

This method creates and returns a new copy of the TensorLike instance on the CPU.

Returns:

A copy of the TensorLike instance placed on the CPU.

Return type:

TensorLike

classmethod from_numbers(numbers, par, batch_mode=None, move_to_numbers_device=True)[source]

Construct an index helper instance from atomic numbers and a parametrization.

Note that this always runs on CPU to avoid inefficient communication between devices. Only the resulting tensors are transfered to the GPU. This is necessary because of complex data look up that is not vectorizable and requires native for-loops. Furthermore, the method frequently uses the torch.Tensor.item() method, which forces CPU-GPU synchronization because it converts a GPU tensor to a Python scalar.

Parameters:

  • numbers (Tensor) – Atomic numbers for all atoms in the system (shape: (..., nat)).

  • par (Param) – Representation of an extended tight-binding model.

  • batch_mode (int) –

    Whether multiple systems or a single one are handled:

    • 0: Single system

    • 1: Multiple systems with padding

    • 2: Multiple systems with no padding (conformer ensemble)

  • move_to_numbers_device (bool) – Move the resulting tensors to the device of the numbers tensor. This should be switched off for GPU calculations that use libcint for integrals as the IndexHelper has to be on the CPU for this step.

Returns:

Instance of index helper for given basis set.

Return type:

IndexHelper

classmethod from_numbers_angular(numbers, angular, batch_mode=None, move_to_numbers_device=True)[source]

Construct an index helper instance from atomic numbers and their angular momenta. If you are not sure about the angular momenta, use IndexHelper.from_numbers() instead, which simply takes a parametrization.

Note that this always runs on CPU to avoid inefficient communication between devices. Only the resulting tensors are transfered to the GPU. This is necessary because of complex data look up that is not vectorizable and requires native for-loops. Furthermore, the method frequently uses the torch.Tensor.item() method, which forces CPU-GPU synchronization because it converts a GPU tensor to a Python scalar.

Parameters:

  • numbers (Tensor) – Atomic numbers for all atoms in the system (shape: (..., nat)).

  • angular (dict[int, Tensor]) – Map between atomic numbers and angular momenta of all shells.

  • batch_mode (int) –

    Whether multiple systems or a single one are handled:

    • 0: Single system

    • 1: Multiple systems with padding

    • 2: Multiple systems with no padding (conformer ensemble)

  • move_to_numbers_device (bool) – Move the resulting tensors to the device of the numbers tensor. This should be switched off for GPU calculations that use libcint for integrals as the IndexHelper has to be on the CPU for this step.

Returns:

Instance of index helper for given basis set.

Return type:

IndexHelper

get_orbital_indices(shell_idx)[source]

Get orbital indices belong to given shell.

Parameters:

shell_idx (int) – Index of given shell.

Returns:

Index list of orbitals belonging to given shell.

Return type:

Tensor

get_shell_indices(atom_idx)[source]

Get shell indices belong to given atom.

Parameters:

atom_idx (int) – Index of given atom.

Returns:

Index list of shells belonging to given atom.

Return type:

Tensor

orbital_atom_mapping(idx)[source]

Mapping of atom index to orbital index, i.e., return indices of orbitals belonging to given atom. The orbital order is given by IndexHelper.orbitals_to_shell().

Parameters:

idx (int) – Index of target atom.

Returns:

1d-Tensor containing the indices of the orbitals.

Return type:

Tensor

override_device(device)

Override the device of the class object.

Warning

This does not change the device of the underlying tensors. It only changes the device of the class object. Use with caution.

Parameters:

device (torch.device) – Device to override the current device.

Return type:

None

override_dtype(dtype)

Override the dtype of the class object.

Warning

This does not change the dtype of the underlying tensors. It only changes the dtype of the class object. Use with caution.

Parameters:

dtype (torch.dtype) – Floating point dtype to override the current dtype.

Return type:

None

reduce_orbital_to_atom(x, dim=-1, reduce='sum', extra=False)[source]

Reduce orbital-resolved tensor to atom-resolved tensor.

Parameters:

  • x (Tensor) – Orbital-resolved tensor.

  • dim (int | (int, int)) – Dimension to reduce over, defaults to -1.

  • reduce (str) – Reduction method, defaults to “sum”.

  • extra (bool) – Tensor to reduce contains a extra dimension of arbitrary size. Defaults to False.

Returns:

Atom-resolved tensor.

Return type:

Tensor

reduce_orbital_to_shell(x, dim=-1, reduce='sum', extra=False)[source]

Reduce orbital-resolved tensor to shell-resolved tensor.

Parameters:

  • x (Tensor) – Orbital-resolved tensor.

  • dim (int | (int, int)) – Dimension to reduce over, defaults to -1.

  • reduce (str) – Reduction method, defaults to “sum”.

  • extra (bool) – Tensor to reduce contains a extra dimension of arbitrary size. Defaults to False.

Returns:

Shell-resolved tensor.

Return type:

Tensor

reduce_shell_to_atom(x, dim=-1, reduce='sum', extra=False)[source]

Reduce shell-resolved tensor to atom-resolved tensor.

Parameters:

  • x (Tensor) – Shell-resolved tensor

  • dim (int | (int, int)) – Dimension to reduce over, defaults to -1.

  • reduce (str) – Reduction method, defaults to “sum”.

  • extra (bool) – Tensor to reduce contains a extra dimension of arbitrary size. Defaults to False.

Returns:

Atom-resolved tensor.

Return type:

Tensor

restore()[source]

Restore the original index helper after culling.

Return type:

None

spread_atom_to_orbital(x, dim=-1, extra=False)[source]

Spread atom-resolved tensor to orbital-resolved tensor.

Parameters:

  • x (Tensor) – Atom-resolved tensor.

  • dim (int | (int, int)) – Dimension to spread over, defaults to -1.

  • extra (bool) – Tensor to reduce contains a extra dimension of arbitrary size. Defaults to False.

Returns:

Orbital-resolved tensor.

Return type:

Tensor

spread_atom_to_orbital_cart(x, dim=-1, extra=False)[source]

Spread atom-resolved tensor to orbital-resolved tensor.

Parameters:

  • x (Tensor) – Atom-resolved tensor.

  • dim (int | (int, int)) – Dimension to spread over, defaults to -1.

  • extra (bool) – Tensor to reduce contains a extra dimension of arbitrary size. Defaults to False.

Returns:

Orbital-resolved tensor.

Return type:

Tensor

spread_atom_to_shell(x, dim=-1, extra=False)[source]

Spread atom-resolved tensor to shell-resolved tensor.

Parameters:

  • x (Tensor) – Atom-resolved tensor.

  • dim (int | (int, int)) – Dimension to spread over, defaults to -1.

  • extra (bool) – Tensor to reduce contains a extra dimension of arbitrary size. Defaults to False.

Returns:

Shell-resolved tensor.

Return type:

Tensor

spread_shell_to_orbital(x, dim=-1, extra=False)[source]

Spread shell-resolved tensor to orbital-resolved tensor.

Parameters:

  • x (Tensor) – Shell-resolved tensor.

  • dim (int | (int, int)) – Dimension to spread over, defaults to -1.

  • extra (bool) – Tensor to reduce contains a extra dimension of arbitrary size. Defaults to False.

Returns:

Orbital-resolved tensor.

Return type:

Tensor

spread_shell_to_orbital_cart(x, dim=-1, extra=False)[source]

Spread shell-resolved tensor to orbital-resolved tensor.

Parameters:

  • x (Tensor) – Shell-resolved tensor.

  • dim (int | (int, int)) – Dimension to spread over, defaults to -1.

  • extra (bool) – Tensor to reduce contains a extra dimension of arbitrary size. Defaults to False.

Returns:

Orbital-resolved tensor.

Return type:

Tensor

spread_ushell_to_orbital(x, dim=-1, extra=False)[source]

Spread unique shell tensor to orbital-resolved tensor.

Parameters:

  • x (Tensor) – Unique shell tensor.

  • dim (int | (int, int)) – Dimension to spread over, defaults to -1.

  • extra (bool) – Tensor to reduce contains a extra dimension of arbitrary size. Defaults to False.

Returns:

Orbital-resolved tensor.

Return type:

Tensor

spread_ushell_to_orbital_cart(x, dim=-1, extra=False)[source]

Spread unique shell tensor to orbital-resolved tensor.

Parameters:

  • x (Tensor) – Unique shell tensor.

  • dim (int | (int, int)) – Dimension to spread over, defaults to -1.

  • extra (bool) – Tensor to reduce contains a extra dimension of arbitrary size. Defaults to False.

Returns:

Orbital-resolved tensor.

Return type:

Tensor

spread_ushell_to_shell(x, dim=-1, extra=False)[source]

Spread unique shell tensor to shell-resolved tensor.

Parameters:

  • x (Tensor) – Unique shell tensor.

  • dim (int | (int, int)) – Dimension to spread over, defaults to -1.

  • extra (bool) – Tensor to reduce contains a extra dimension of arbitrary size. Defaults to False.

Returns:

Shell-resolved tensor.

Return type:

Tensor

spread_uspecies_to_atom(x, dim=-1, extra=False)[source]

Spread unique species tensor to atom-resolved tensor.

Parameters:

  • x (Tensor) – Unique specie tensor.

  • dim (int | (int, int)) – Dimension to spread over, defaults to -1.

  • extra (bool) – Tensor to reduce contains a extra dimension of arbitrary size. Defaults to False.

Returns:

Atom-resolved tensor.

Return type:

Tensor

spread_uspecies_to_orbital(x, dim=-1, extra=False)[source]

Spread unique species tensor to orbital-resolved tensor.

Parameters:

  • x (Tensor) – Unique specie tensor.

  • dim (int) – Dimension to spread over, defaults to -1.

  • extra (bool) – Tensor to reduce contains a extra dimension of arbitrary size. Defaults to False.

Returns:

Orbital-resolved tensor.

Return type:

Tensor

spread_uspecies_to_orbital_cart(x, dim=-1, extra=False)[source]

Spread unique species tensor to orbital-resolved tensor.

Parameters:

  • x (Tensor) – Unique specie tensor.

  • dim (int) – Dimension to spread over, defaults to -1.

  • extra (bool) – Tensor to reduce contains a extra dimension of arbitrary size. Defaults to False.

Returns:

Orbital-resolved tensor.

Return type:

Tensor

spread_uspecies_to_shell(x, dim=-1, extra=False)[source]

Spread unique species tensor to shell-resolved tensor.

Parameters:

  • x (Tensor) – Unique specie tensor.

  • dim (int | (int, int)) – Dimension to spread over, defaults to -1.

  • extra (bool) – Tensor to reduce contains a extra dimension of arbitrary size. Defaults to False.

Returns:

Shell-resolved tensor.

Return type:

Tensor

to(device=None, dtype=None)

Returns a copy of the TensorLike instance on the specified device.

This method creates and returns a new copy of the TensorLike instance on the specified device “device”.

Parameters:

  • device (torch.device) – Device to which all associated tensors should be moved.

  • dtype (dtype | None)

Returns:

A copy of the TensorLike instance placed on the specified device.

Return type:

TensorLike

Notes

If the TensorLike instance is already on the desired device self will be returned.

type(dtype)

Returns a copy of the TensorLike instance with specified floating point type. This method creates and returns a new copy of the TensorLike instance with the specified dtype.

Parameters:

dtype (torch.dtype) – Floating point type.

Returns:

A copy of the TensorLike instance with the specified dtype.

Return type:

TensorLike

Notes

If the TensorLike instance has already the desired dtype Self will be returned.

property allowed_dtypes: tuple[dtype, ...]

Specification of dtypes that the TensorLike object can take.

Returns:

Collection of allowed dtypes the TensorLike object can take.

Return type:

tuple[torch.dtype, …]

angular: Tensor

Angular momenta for all shells

atom_to_unique: Tensor

Mapping of atoms to unique species

batch_mode: int

Whether multiple systems or a single one are handled:

  • 0: Single system

  • 1: Multiple systems with padding

  • 2: Multiple systems with no padding (conformer ensemble)

property dd: DD

Shortcut for device and dtype.

property device: device

The device on which the class object resides.

property dtype: dtype

Floating point dtype used by class object.

orbital_index: Tensor

Offset index for starting the next orbital block

property orbitals_per_atom: Tensor

Number of orbitals for each atom.

Returns:

Atom indices for each orbital.

Return type:

Tensor

orbitals_per_shell: Tensor

Number of orbitals for each shell

orbitals_to_shell: Tensor

Mapping of orbitals to shells

shell_index: Tensor

Offset index for starting the next shell block

shells_per_atom: Tensor

Number of shells for each atom

shells_to_atom: Tensor

Mapping of shells to atoms

shells_to_ushell: Tensor

Mapping of shells to unique atoms

store: IndexHelperStore | None

Storage to restore from after culling.

unique_angular: Tensor

Angular momenta of all unique shells

ushells_per_unique: Tensor

Number of unique shells per unqiue atoms.

ushells_to_unique: Tensor

Mapping of unique shells to unique species

Basis Construction#

Basis: Main Class#

Main basis set class for creating the contracted Gaussian type orbitals (CGTOs) from the parametrization. The basis set can also be printed in various formats.

class dxtb._src.basis.bas.Basis(numbers, par, ihelp, device=None, dtype=None)[source]#

Bases: TensorLike

Atomic orbital basis set.

Parameters:
cpu()#

Returns a copy of the TensorLike instance on the CPU.

This method creates and returns a new copy of the TensorLike instance on the CPU.

Returns:

A copy of the TensorLike instance placed on the CPU.

Return type:

TensorLike

create_cgtos()[source]#

Create contracted Gaussian type orbitals from parametrization.

Returns:

List of primitive Gaussian exponents and contraction coefficients for the orthonormalized basis functions for each shell.

Return type:

tuple[list[Tensor], list[Tensor]]

create_libcint(positions, mask=None)[source]#

Create the basis set required for libcint.

Parameters:

  • positions (Tensor) – Cartesian coordinates of all atoms (shape: (..., nat, 3)).

  • mask (Tensor | None, optional) – Mask for positions to make batched computations easier. The overlap does not work in a batched fashion. Hence, we loop over the batch dimension and must remove the padding. Defaults to None, i.e., tad_mctc.batch.deflate() is used.

Returns:

List of CGTOs.

Return type:

list[libcint.AtomCGTOBasis] | list[list[libcint.AtomCGTOBasis]]

Raises:

NotImplementedError – If batch mode is requested (checked through dimensions of numbers).

override_device(device)#

Override the device of the class object.

Warning

This does not change the device of the underlying tensors. It only changes the device of the class object. Use with caution.

Parameters:

device (torch.device) – Device to override the current device.

Return type:

None

override_dtype(dtype)#

Override the dtype of the class object.

Warning

This does not change the dtype of the underlying tensors. It only changes the dtype of the class object. Use with caution.

Parameters:

dtype (torch.dtype) – Floating point dtype to override the current dtype.

Return type:

None

to(device=None, dtype=None)#

Returns a copy of the TensorLike instance on the specified device.

This method creates and returns a new copy of the TensorLike instance on the specified device “device”.

Parameters:

  • device (torch.device) – Device to which all associated tensors should be moved.

  • dtype (dtype | None)

Returns:

A copy of the TensorLike instance placed on the specified device.

Return type:

TensorLike

Notes

If the TensorLike instance is already on the desired device self will be returned.

to_bse(qcformat='nwchem', save=False, overwrite=False, verbose=False, with_header=False)[source]#

Convert the basis set to a format suitable for basis set exchange.

Parameters:

  • qcformat (Literal["gaussian94", "nwchem"], optional) – Format of the basis set. Defaults to "nwchem".

  • save (bool, optional) – Whether to save the basis set to a file. Defaults to False.

  • overwrite (bool, optional) – Whether to overwrite existing files. Defaults to False.

  • verbose (bool, optional) – Whether to print the basis set to the console. Defaults to False.

  • with_header (bool, optional) – Whether to include the header in the basis set. Defaults to False.

Returns:

Basis set in the specified format.

Return type:

str

Raises:

  • RuntimeError – If no meta data is found in the parametrization, or if the meta data is incomplete (name missing), or if the basis set is batched.

  • ValueError – If no atoms for basis set printout are found, or if the basis set format is not supported.

type(dtype)#

Returns a copy of the TensorLike instance with specified floating point type. This method creates and returns a new copy of the TensorLike instance with the specified dtype.

Parameters:

dtype (torch.dtype) – Floating point type.

Returns:

A copy of the TensorLike instance with the specified dtype.

Return type:

TensorLike

Notes

If the TensorLike instance has already the desired dtype Self will be returned.

property allowed_dtypes: tuple[dtype, ...]#

Specification of dtypes that the TensorLike object can take. Defaults to float types and must be overridden by subclass if float are not allowed. The IndexHelper is an example that should only allow integers.

Returns:

Collection of allowed dtypes the TensorLike object can take.

Return type:

tuple[torch.dtype, …]

property dd: DD#

Shortcut for device and dtype.

property device: device#

The device on which the class object resides.

property dtype: dtype#

Floating point dtype used by class object.

ngauss: Tensor#

Number of Gaussians used in expansion from Slater orbital.

pqn: Tensor#

Principal quantum number of each shell

shells: dict[str, list[str]]#

Shells for each atom.

slater: Tensor#

Exponent of Slater function.

valence: Tensor#

Whether the shell is part of the valence shell.

Basis: Orthonormalization#

Gram-Schmidt orthonormalization routines for contracted Gaussian basis functions.

dxtb._src.basis.ortho.gaussian_integral(ai, aj, ci, cj)[source]#

Integral over two Gaussians (overlap).

Parameters:

  • ai (Tensor) – Exponent of GTO i. Can also be a tensor of exponents corresponding to all primitive GTOs of GTO i.

  • aj (Tensor) – Exponent of GTO j. Can also be a tensor of exponents corresponding to all primitive GTOs of GTO j.

  • ci (Tensor) – Contraction coefficients for CGTO i.

  • cj (Tensor) – Contraction coefficients for CGTO j.

Returns:

Overlap (summed).

Return type:

Tensor

dxtb._src.basis.ortho.orthogonalize(alpha, coeff)[source]#

Orthogonalize a contracted Gaussian basis function to an existing basis function using. The second basis function is orthonormalized against the first basis function.

Parameters:

  • alpha ((Tensor, Tensor)) – Primitive Gaussian exponents for the shell pair.

  • coeff ((Tensor, Tensor)) – Contraction coefficients for the shell pair.

Returns:

Primitive Gaussian exponents and contraction coefficients for the orthonormalized basis function.

Return type:

(Tensor, Tensor)

Basis: Slater Expansion#

Expansion coefficients for Slater functions into primitive Gaussian functions

dxtb._src.basis.slater.slater_to_gauss(ng, n, l, zeta, norm=True)[source]#

Expand Slater function in primitive gaussian functions.

Parameters:

  • ng (int) – Number of Gaussian functions for the expansion.

  • n (int) – Principal quantum number of shell.

  • l (int) – Azimuthal quantum number of shell.

  • zeta (Tensor) – Exponent of Slater function to expand.

  • norm (bool, optional) – Include normalization in contraction coefficients. Defaults to True.

Returns:

Contraction coefficients of primitive gaussians, can contain normalization, and exponents of primitive gaussian functions.

Return type:

(Tensor, Tensor)

Analysis#

Wavefunction#

Provides methods to create and analyze wavefunctions.

Wavefunction: Filling#

Handle the occupation of the orbitals with electrons.

Parts of the Fermi smearing are taken from tbmalt/tbmalt. The Fermi energy search follows the Fermi filling of tblite after pull request #385 (tblite/tblite#385, commit e437cde), with additions for batches, fractional electrons and derivatives (see get_fermi_occupation()). The derivatives of the occupations are obtained by differentiating Newton steps after a detached solve, an instance of one-step differentiation [Bolte2023] that is extended here to higher orders.

References

[Bolte2023] (1,2)

J. Bolte, E. Pauwels, S. Vaiter. One-step differentiation of iterative algorithms. 37th Conference on Neural Information Processing Systems (NeurIPS 2023). Corollary 2 (vanishing Jacobian at the fixed point) and Corollary 3 (Newton’s method, quadratic convergence) prove the first-order result. The paper contains no statement about higher orders.

dxtb._src.wavefunction.filling.get_alpha_beta_occupation(nel, uhf=None)[source]#

Generate alpha and beta electrons from total number of electrons.

Parameters:

  • nel (Tensor) – Total number of electrons.

  • uhf (Tensor | int | list[int] | None) – Number of unpaired electrons. If None, spin is figured out automatically.

Returns:

Alpha (first column, 0 index) and beta (second column, 1 index) electrons.

Return type:

Tensor

Raises:

ValueError – Number of electrons and unpaired electrons does not match.

Note

Rounding is only used to determine the parity of the number of electrons, which fixes the number of unpaired electrons. The number of electrons itself is not rounded: the paired electrons of a fractional number are split equally between alpha and beta, the unpaired ones are alpha, as in tblite. At half-integer totals, the parity and hence the split change discontinuously.

dxtb._src.wavefunction.filling.get_aufbau_occupation(norb, nel)[source]#

Set occupation numbers according to the aufbau principle. The number of electrons is a real number and can be fractional. Orbitals beyond the number of available orbitals norb of a system (padding in a batch) are never occupied.

Parameters:

  • norb (Tensor) – Number of available orbitals.

  • nel (Tensor) – Number of electrons.

Returns:

Occupation numbers.

Return type:

Tensor

Examples

>>> get_aufbau_occupation(torch.tensor(5), torch.tensor(1.))
tensor([1., 0., 0., 0., 0.])
>>> get_aufbau_occupation(torch.tensor([8, 8, 5]), torch.tensor([2., 3., 1.]))
tensor([[1., 1., 0., 0., 0., 0., 0., 0.],
        [1., 1., 1., 0., 0., 0., 0., 0.],
        [1., 0., 0., 0., 0., 0., 0., 0.]])
>>> nel, norb = torch.tensor([2.0, 3.5, 1.5]), torch.tensor([4, 4, 2])
>>> occ = get_aufbau_occupation(norb, nel)
>>> occ
tensor([[1.0000, 1.0000, 0.0000, 0.0000],
        [1.0000, 1.0000, 1.0000, 0.5000],
        [1.0000, 0.5000, 0.0000, 0.0000]])
>>> all(nel == occ.sum(-1))
True

import torch
from dxtb.wavefunction import get_aufbau_occupation

# 1 electron in 5 orbitals
r1 = get_aufbau_occupation(torch.tensor(5), torch.tensor(1.))

print(r1)
# Output: tensor([1., 0., 0., 0., 0.])

# Multiple orbitals and different electron counts
r2 = get_aufbau_occupation(
    torch.tensor([8, 8, 5]), torch.tensor([2., 3., 1.])
)

print(r2)
# Output: tensor([[1., 1., 0., 0., 0., 0., 0., 0.],
#                 [1., 1., 1., 0., 0., 0., 0., 0.],
#                 [1., 0., 0., 0., 0., 0., 0., 0.]])

# Fractional electron numbers in multiple orbitals
nel, norb = torch.tensor([2.0, 3.5, 1.5]), torch.tensor([4, 4, 2])
occ = get_aufbau_occupation(norb, nel)

print(occ)
# Output: tensor([[1.0000, 1.0000, 0.0000, 0.0000],
#                 [1.0000, 1.0000, 1.0000, 0.5000],
#                 [1.0000, 0.5000, 0.0000, 0.0000]])

# Check if the total number of electrons matches the sum of occupation
print(all(nel == occ.sum(-1)))  # True
dxtb._src.wavefunction.filling.get_fermi_energy(nel, emo, mask=None)[source]#

Get Fermi energy as midpoint between the HOMO and LUMO.

The orbital energies emo and the mask must already have the correct shape for using alpha/beta electron channels. Spreading to the channels can be done with x.unsqueeze(-2).expand([*nel.shape, -1]).

Parameters:

  • nel (Tensor) – Number of electrons per channel (shape [b, 2], the batch dimension b is optional).

  • emo (Tensor) – Orbital energies (shape [b, 2, n], the same for both channels).

  • mask (Tensor | None, optional) – Mask from orbitals to avoid reading padding as LUMO for elements without LUMO due to minimal basis (shape [b, 2, n]).

Returns:

Fermi energy (shape [b, 2]) and index of HOMO (shape [b, 2, 1]).

Return type:

tuple[Tensor, Tensor]

dxtb._src.wavefunction.filling.get_fermi_occupation(nel, emo, kt, mask=None, thr=None, maxiter=200, diff_order=None)[source]#

Set occupation numbers according to Fermi distribution.

The Fermi energy is determined such that the occupations sum up to the given (possibly fractional) number of electrons nel. The algorithm is the one of tblite (Newton iteration on the number of electrons) with three additions that are necessary for a batched, differentiable implementation:

  1. The Newton iteration is safeguarded by a bisection bracket because it is not globally convergent, e.g., for fractional electrons. Converged entries of a batch are frozen. This part does not track gradients.

  2. The converged Fermi energy is polished by Newton steps on the residual in log space, which pins it to machine precision also in a gap, where any Fermi energy between the orbitals converges the number of electrons. This part does not track gradients either.

  3. Differentiable Newton steps from the converged Fermi energy carry its derivatives up to the requested order diff_order (see the note on derivatives below) and are attached to the graph afterwards.

The orbital energies emo must already have the correct shape for using alpha/beta electron channels. Spreading to the channels can be done with emo.unsqueeze(-2).expand([*nel.shape, -1]).

Shapes (the batch dimension b is optional, n is the number of orbitals including padding): the electrons nel are [b, 2], the orbital energies emo, the mask and the occupations are [b, 2, n]. Each channel (alpha, beta) is optimized independently and the orbitals are at most singly occupied. Intermediate quantities per channel, such as the Fermi energy or the number of electrons, have a trailing singleton dimension ([b, 2, 1]) to broadcast against the orbitals.

Parameters:

  • nel (Tensor) – Number of electrons per channel ([b, 2]). It may be fractional, and the occupations are differentiable with respect to it (e.g., for the chemical potential or Fukui functions). A graph that must not be part of the occupations (e.g., the one of the previous SCF step) has to be detached by the caller.

  • emo (Tensor) – Orbital energies ([b, 2, n]).

  • kt (Tensor) – Electronic temperature in atomic units (scalar). It is moved to the device of emo. For kt == 0, the aufbau occupation is returned.

  • mask (Tensor | None, optional) – Mask for the existing orbitals (0 for padding) with the same shape as emo. Padded orbitals are never occupied and are not read as LUMO for the initial guess of the Fermi energy. Without a mask, padding cannot be distinguished from actual orbitals and is occupied if its energy is close to the Fermi energy.

  • thr (Tensor | float | int | None, optional) – Threshold for the deviation of the number of electrons, by default None, which is min(sqrt(eps), 1e5 * eps, 1e-4) of the dtype of emo (the last limit only applies in single precision). The SCF verifies the number of electrons to about 5e-4, i.e., larger thresholds may trip this check in single precision.

  • maxiter (int, optional) – Maximum number of iterations for converging Fermi energy. Defaults to 200.

  • diff_order (int | None, optional) – Highest order of the derivatives of the occupations that is exact (with respect to the orbital energies, the temperature, the number of electrons and everything upstream). k = ceil(log2(diff_order + 1)) differentiable Newton steps are attached, i.e., 0 for order 0, 1 for order 1, 2 for orders 2 and 3, and 3 for orders 4 to 7. The default None is FERMI_DIFF_ORDER (order 3: forces, Hessians, polarizabilities, dipole derivatives and first hyperpolarizabilities), set with the fermi_diff_order option in the SCF. The value of the occupations does not depend on it (apart from the forward residual, see below). Each additional order of the derivative costs more than the additional Newton steps, since the nested derivatives of the graph grow exponentially. Orders beyond 3 need double precision.

Returns:

Occupation numbers.

Return type:

Tensor

Raises:

  • RuntimeError – Fermi energy fails to converge.

  • TypeError – Electronic temperature is not given as Tensor or the derivative order is not an integer.

  • ValueError – Electronic temperature is not a scalar or negative, the number of electrons exceeds the number of orbitals, or the derivative order is negative.

Note

Derivatives (with respect to the orbital energies, and everything upstream of them, such as positions and fields, to kt and to nel):

  • k undamped Newton steps mu <- mu - g / g' of the number of electrons g from a converged, detached start reproduce all derivatives of the exact Fermi energy up to order 2**k - 1. This is the reason for k = ceil(log2(diff_order + 1)). Proof: the Newton map N fulfils N(mu*) = mu* and N'(mu*) = 0, i.e., N(mu) - mu* = (mu - mu*)**2 h with a smooth h (roughly g'' / (2 g')). With the start mu_0 = mu*(eps_0), the deviation mu_0 - mu*(eps) is O(d eps), and applying the identity k times gives mu_k - mu* = O(d eps**(2**k)). By the chain rule, the derivatives of the occupations agree up to the same order. The argument holds for any parameter of the Newton map, i.e., also for the number of electrons (d mu / d N = 1 / g') and for mixed derivatives.

  • The derivative of the occupations of an active channel with respect to its number of electrons sums to one over the orbitals. Channels without electrons and completely filled channels have no derivative with respect to it (they are not optimized). At zero temperature, it is the derivative of the aufbau filling: one for the partially occupied orbital, and for integer electrons for the lowest empty one (the derivative for adding electrons).

  • In a gap (integer electrons), an additional electron is distributed according to the thermal weights, d f_i / d N = w_i / sum_j w_j with w = f (1 - f), i.e., mostly to the HOMO and the LUMO, which are tails of the Fermi function. Three parts keep this ratio exact far beyond the threshold of the search: the residual of the Newton steps is evaluated without cancellation (holes below and electrons above the integer number of electrons, see _partition), the Fermi energy is polished in log space (its position in the gap sets the ratio), and the Fermi function has no cutoff.

  • The derivatives are exact up to the forward residual e_0 of the Fermi energy: the n-th derivative has an error of the order of e_0**(2**k - n). For diff_order = 2**k - 1, the highest order has an error of order e_0 (the error of a hand-written implicit derivative), and the lower orders are more accurate.

  • The default order 3 (two steps) covers forces, Hessians, polarizabilities, dipole derivatives (up to second order in the occupations) and the first hyperpolarizability and derivative of the polarizability (third order). Fourth-order properties need order 4 or more (three steps).

  • Order 0 (no step) omits the change of the Fermi energy entirely, which makes even forces wrong. Order 1 (one step) is exact for forces only.

  • The origin of the idea is one-step differentiation (Bolte, Pauwels and Vaiter, NeurIPS 2023, see the references of this module, [Bolte2023]): a detached solve, differentiate only the last step(s) of a fast algorithm. Corollary 2 shows that one step gives the exact Jacobian if the Jacobian of the iteration map vanishes at the fixed point, and Corollary 3 bounds the error of one step by L_J ||x_{k-1} - x*|| for quadratically convergent maps such as Newton’s method. Both are first-order results. The statement for higher orders (2**k - 1) is the proof above and not part of the paper.

  • The steps must move the value. A straight-through form mu + (change - change.detach()) differentiates the occupations at the start instead of the end of the steps and leaves an error of order e_0 for every k.

  • Completely filled channels (as many electrons as orbitals, e.g., He or H-) have no finite Fermi energy. Their occupations are exactly one and constant, i.e., all derivatives vanish.

  • The derivative of the Fermi energy is dropped (the Fermi energy stays frozen) if the derivative of the number of electrons, g' = sum f (1 - f) / kT, is below (tiny / eps)**(1 / 2**k), or if the step is longer than _MAX_DIFF_STEP_KT times kT (a safety net, not reached after the polish). Below the floor, the derivatives of order 2**k with respect to nel would overflow, since the m-th one grows like g'**(1 - m) in a gap. The missing terms for the orbital energies and kt are of the order of f (1 - f), i.e., negligible. For nel, the derivative of that channel is zero instead of summing to one over the orbitals. With the Fermi energy in the middle of a gap and kT = 1e-3, this happens for gaps of more than about 690, 355 and 180 kT for the orders 1, 3 and 7 in double precision and 87 and 55 kT for the orders 1 and 3 in single precision (measured). Up to order 2**k (one beyond the exact ones), the derivatives stay finite (measured); higher orders may overflow in a gap, since their true values do.

  • The search reads its convergence flags on the host, i.e., the function has data-dependent control flow. Autograd and the transforms jacrev, jacfwd, hessian and jvp of torch.func work, vmap does not, and torch.compile breaks the graph at every convergence check.

Wavefunction: Mulliken#

Wavefunction analysis via Mulliken populations.

dxtb._src.wavefunction.mulliken.get_atomic_populations(overlap, density, indexhelper)[source]#

Compute atom-resolved populations.

Parameters:

  • overlap (Tensor) – Overlap matrix.

  • density (Tensor) – Density matrix.

  • indexhelper (IndexHelper) – Index mapping for the basis set.

Returns:

Atom populations.

Return type:

Tensor

dxtb._src.wavefunction.mulliken.get_mulliken_atomic_charges(overlap, density, indexhelper, n0)[source]#

Compute atom-resolved Mulliken partial charges.

Parameters:

  • overlap (Tensor) – Overlap matrix.

  • density (Tensor) – Density matrix.

  • indexhelper (IndexHelper) – Index mapping for the basis set.

  • n0 (Tensor) – Atom-resolved reference occupancy numbers.

Returns:

Atom-resolved Mulliken partial charges.

Return type:

Tensor

dxtb._src.wavefunction.mulliken.get_mulliken_shell_charges(overlap, density, indexhelper, n0)[source]#

Compute shell-resolved Mulliken partial charges using Mulliken population analysis.

Parameters:

  • overlap (Tensor) – Overlap matrix.

  • density (Tensor) – Density matrix.

  • indexhelper (IndexHelper) – Index mapping for the basis set.

  • n0 (Tensor) – Shell-resolved reference occupancy numbers.

Returns:

Shell-resolved Mulliken partial charges.

Return type:

Tensor

dxtb._src.wavefunction.mulliken.get_orbital_populations(overlap, density)[source]#

Compute orbital-resolved populations using Mulliken population analysis.

Parameters:

  • overlap (Tensor) – Overlap matrix.

  • density (Tensor) – Density matrix.

Returns:

Orbital populations.

Return type:

Tensor

dxtb._src.wavefunction.mulliken.get_shell_populations(overlap, density, indexhelper)[source]#

Compute shell-resolved populations using Mulliken population analysis.

Parameters:

  • overlap (Tensor) – Overlap matrix.

  • density (Tensor) – Density matrix.

  • indexhelper (IndexHelper) – Index mapping for the basis set.

Returns:

Shell populations.

Return type:

Tensor

Wavefunction: Wiberg/Mayer Bond Orders#

Wiberg (or better Mayer) bond orders are calculated from the off-diagonal elements of the matrix product of the density and the overlap matrix.

dxtb._src.wavefunction.wiberg.get_bond_order(overlap, density, ihelp)[source]#

Calculate Wiberg bond orders.

Parameters:

  • overlap (Tensor) – Overlap matrix.

  • density (Tensor) – Density matrix.

  • ihelp (IndexHelper) – Helper class for indexing.

Returns:

Wiberg bond orders.

Return type:

Tensor

Contents