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.

Return type:

Self

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.

Return type:

IndexHelper

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

Parameters:
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.

Return type:

IndexHelper

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

Parameters:
get_orbital_indices(shell_idx)[source]

Get orbital indices belong to given shell.

Return type:

Tensor

Parameters:

shell_idx (int) – Index of given shell.

Returns:

Index list of orbitals belonging to given shell.

Return type:

Tensor

Parameters:

shell_idx (int)

get_shell_indices(atom_idx)[source]

Get shell indices belong to given atom.

Return type:

Tensor

Parameters:

atom_idx (int) – Index of given atom.

Returns:

Index list of shells belonging to given atom.

Return type:

Tensor

Parameters:

atom_idx (int)

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().

Return type:

Tensor

Parameters:

idx (int) – Index of target atom.

Returns:

1d-Tensor containing the indices of the orbitals.

Return type:

Tensor

Parameters:

idx (int)

override_device(device)

Override the device of the class object. :rtype: None

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.

Parameters:

device (device)

Return type:

None

override_dtype(dtype)

Override the dtype of the class object. :rtype: None

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.

Parameters:

dtype (dtype)

Return type:

None

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

Reduce orbital-resolved tensor to atom-resolved tensor.

Return type:

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

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

Reduce orbital-resolved tensor to shell-resolved tensor.

Return type:

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

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

Reduce shell-resolved tensor to atom-resolved tensor.

Return type:

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

Parameters:
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.

Return type:

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

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

Spread atom-resolved tensor to orbital-resolved tensor.

Return type:

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

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

Spread atom-resolved tensor to shell-resolved tensor.

Return type:

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

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

Spread shell-resolved tensor to orbital-resolved tensor.

Return type:

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

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

Spread shell-resolved tensor to orbital-resolved tensor.

Return type:

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

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

Spread unique shell tensor to orbital-resolved tensor.

Return type:

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

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

Spread unique shell tensor to orbital-resolved tensor.

Return type:

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

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

Spread unique shell tensor to shell-resolved tensor.

Return type:

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

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

Spread unique species tensor to atom-resolved tensor.

Return type:

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

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

Spread unique species tensor to orbital-resolved tensor.

Return type:

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

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

Spread unique species tensor to orbital-resolved tensor.

Return type:

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

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

Spread unique species tensor to shell-resolved tensor.

Return type:

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

Parameters:
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”.

Return type:

Self

Parameters:

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

Returns:

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

Return type:

TensorLike

Parameters:

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.

Return type:

Self

Parameters:

dtype (torch.dtype) – Floating point type.

Returns:

A copy of the TensorLike instance with the specified dtype.

Return type:

TensorLike

Parameters:

dtype (dtype)

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.

Return type:

Self

Returns:

A copy of the TensorLike instance placed on the CPU.

Return type:

TensorLike

create_cgtos()[source]#

Create contracted Gaussian type orbitals from parametrization.

Return type:

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

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.

Return type:

list[AtomCGTOBasis] | list[list[AtomCGTOBasis]]

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).

Parameters:
override_device(device)#

Override the device of the class object. :rtype: None

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.

Parameters:

device (device)

Return type:

None

override_dtype(dtype)#

Override the dtype of the class object. :rtype: None

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.

Parameters:

dtype (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”.

Return type:

Self

Parameters:

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

Returns:

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

Return type:

TensorLike

Parameters:

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.

Return type:

str

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.

Parameters:
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.

Return type:

Self

Parameters:

dtype (torch.dtype) – Floating point type.

Returns:

A copy of the TensorLike instance with the specified dtype.

Return type:

TensorLike

Parameters:

dtype (dtype)

Notes

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

unique_shell_pairs(mask=None, uplo='l')[source]#

Create a matrix of unique shell pairs.

Return type:

tuple[Tensor, Tensor]

Parameters:
  • mask (Tensor) – Mask for

  • uplo (Literal[“n”, “N”, “u”, “U”, “l”, “L”]) – Whether the matrix of unique shell pairs should be create as a triangular matrix (l: lower, u: upper) or full matrix (n). Defaults to l (lower triangular matrix).

Returns:

Matrix of unique shell pairs and the number of unique shell pairs.

Return type:

(Tensor, Tensor)

Raises:

ValueError – The number of unique shell pairs does not match the theoretical one.

Parameters:
  • mask (Tensor | None)

  • uplo (Literal['n', 'N', 'u', 'U', 'l', 'L'])

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).

Return type:

Tensor

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

Parameters:
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.

Return type:

tuple[Tensor, Tensor]

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)

Parameters:

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.

Return type:

tuple[Tensor, Tensor]

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)

Parameters:

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

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

Generate alpha and beta electrons from total number of electrons.

Return type:

Tensor

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.

Parameters:

Note

The number of electrons is rounded to integers via torch.round for numerical stability, i.e., non-integer electrons are not supported.

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.

Return type:

Tensor

Parameters:
  • norb (Tensor) – Number of available orbitals.

  • nel (Tensor) – Number of electrons.

Returns:

Occupation numbers.

Return type:

Tensor

Parameters:

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]).

Return type:

tuple[Tensor, Tensor]

Parameters:
  • nel (Tensor) – Number of electrons.

  • emo (Tensor) – Orbital energies

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

Returns:

Fermi energy and index of HOMO.

Return type:

tuple[Tensor, Tensor]

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

Set occupation numbers according to Fermi distribution.

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]).

Return type:

Tensor

Parameters:
  • nel (Tensor) – Number of electrons.

  • emo (Tensor) – Orbital energies.

  • kt (Tensor) – Electronic temperature in atomic units.

  • mask (Tensor | None, optional) – Mask for Fermi energy. Just passed through.

  • thr (Tensor | None, optional) – Threshold for converging Fermi energy, by default None.

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

Returns:

Occupation numbers.

Return type:

Tensor

Raises:
  • RuntimeError – Fermi energy fails to converge.

  • TypeError – Electronic temperature is not given as Tensor.

  • ValueError – Electronic temperature is negative or number of electrons is zero.

Parameters:

Wavefunction: Mulliken#

Wavefunction analysis via Mulliken populations.

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

Compute atom-resolved populations.

Return type:

Tensor

Parameters:
  • overlap (Tensor) – Overlap matrix.

  • density (Tensor) – Density matrix.

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

Returns:

Atom populations.

Return type:

Tensor

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

Compute atom-resolved Mulliken partial charges.

Return type:

Tensor

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

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

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

Return type:

Tensor

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

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

Compute orbital-resolved populations using Mulliken population analysis.

Return type:

Tensor

Parameters:
  • overlap (Tensor) – Overlap matrix.

  • density (Tensor) – Density matrix.

Returns:

Orbital populations.

Return type:

Tensor

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

Compute shell-resolved populations using Mulliken population analysis.

Return type:

Tensor

Parameters:
  • overlap (Tensor) – Overlap matrix.

  • density (Tensor) – Density matrix.

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

Returns:

Shell populations.

Return type:

Tensor

Parameters:

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.

Return type:

Tensor

Parameters:
  • overlap (Tensor) – Overlap matrix.

  • density (Tensor) – Density matrix.

  • ihelp (IndexHelper) – Helper class for indexing.

Returns:

Wiberg bond orders.

Return type:

Tensor

Parameters: