dxtb.IndexHelper

dxtb.IndexHelper#

class dxtb.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.

Methods

`cpu`	Returns a copy of the `TensorLike` instance on the CPU.
`cull`
`from_numbers`	Construct an index helper instance from atomic numbers and a parametrization.
`from_numbers_angular`	Construct an index helper instance from atomic numbers and their angular momenta.
`get_orbital_indices`	Get orbital indices belong to given shell.
`get_shell_indices`	Get shell indices belong to given atom.
`orbital_atom_mapping`	Mapping of atom index to orbital index, i.e., return indices of orbitals belonging to given atom.
`override_device`	Override the device of the class object.
`override_dtype`	Override the dtype of the class object.
`reduce_orbital_to_atom`	Reduce orbital-resolved tensor to atom-resolved tensor.
`reduce_orbital_to_shell`	Reduce orbital-resolved tensor to shell-resolved tensor.
`reduce_shell_to_atom`	Reduce shell-resolved tensor to atom-resolved tensor.
`restore`	Restore the original index helper after culling.
`spread_atom_to_orbital`	Spread atom-resolved tensor to orbital-resolved tensor.
`spread_atom_to_orbital_cart`	Spread atom-resolved tensor to orbital-resolved tensor.
`spread_atom_to_shell`	Spread atom-resolved tensor to shell-resolved tensor.
`spread_shell_to_orbital`	Spread shell-resolved tensor to orbital-resolved tensor.
`spread_shell_to_orbital_cart`	Spread shell-resolved tensor to orbital-resolved tensor.
`spread_ushell_to_orbital`	Spread unique shell tensor to orbital-resolved tensor.
`spread_ushell_to_orbital_cart`	Spread unique shell tensor to orbital-resolved tensor.
`spread_ushell_to_shell`	Spread unique shell tensor to shell-resolved tensor.
`spread_uspecies_to_atom`	Spread unique species tensor to atom-resolved tensor.
`spread_uspecies_to_orbital`	Spread unique species tensor to orbital-resolved tensor.
`spread_uspecies_to_orbital_cart`	Spread unique species tensor to orbital-resolved tensor.
`spread_uspecies_to_shell`	Spread unique species tensor to shell-resolved tensor.
`to`	Returns a copy of the `TensorLike` instance on the specified device.
`type`	Returns a copy of the `TensorLike` instance with specified floating point type.

Attributes

`unique_angular`	Angular momenta of all unique shells
`angular`	Angular momenta for all shells
`atom_to_unique`	Mapping of atoms to unique species
`ushells_to_unique`	Mapping of unique shells to unique species
`ushells_per_unique`	Number of unique shells per unqiue atoms.
`shells_to_ushell`	Mapping of shells to unique atoms
`shells_per_atom`	Number of shells for each atom
`shell_index`	Offset index for starting the next shell block
`shells_to_atom`	Mapping of shells to atoms
`orbitals_per_shell`	Number of orbitals for each shell
`orbital_index`	Offset index for starting the next orbital block
`orbitals_to_shell`	Mapping of orbitals to shells
`batch_mode`	Whether multiple systems or a single one are handled:
`store`	Storage to restore from after culling.
`allowed_dtypes`	Specification of dtypes that the TensorLike object can take.
`dd`	Shortcut for device and dtype.
`device`	The device on which the class object resides.
`dtype`	Floating point dtype used by class object.
`nao`
`nat`
`nbatch`
`nsh`
`orbital_index_cart`
`orbitals_per_atom`	Number of orbitals for each atom.
`orbitals_per_shell_cart`
`orbitals_to_atom`
`orbitals_to_atom_cart`
`orbitals_to_shell_cart`

Parameters:

unique_angular (Tensor)
angular (Tensor)
atom_to_unique (Tensor)
ushells_to_unique (Tensor)
ushells_per_unique (Tensor)
shells_to_ushell (Tensor)
shells_per_atom (Tensor)
shell_index (Tensor)
shells_to_atom (Tensor)
orbitals_per_shell (Tensor)
orbital_index (Tensor)
orbitals_to_shell (Tensor)
batch_mode (int)
device (torch.device | None)
dtype (torch.dtype)
store (IndexHelperStore | None)

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

cull(conv, slicers)[source]#

Return type:

None

Parameters:

conv (Tensor)
slicers (Slicers)

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:

numbers (Tensor)
par (Param | ParamModule)
batch_mode (int | None)
move_to_numbers_device (bool)

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:

numbers (Tensor)
angular (dict[int, list[int]])
batch_mode (int | None)
move_to_numbers_device (bool)

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:

x (Tensor)
dim (int | tuple[int, int])
reduce (str)
extra (bool)

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:

x (Tensor)
dim (int | tuple[int, int])
reduce (str)
extra (bool)

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:

x (Tensor)
dim (int | tuple[int, int])
reduce (str)
extra (bool)

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:

x (Tensor)
dim (int | tuple[int, int])
extra (bool)

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:

x (Tensor)
dim (int | tuple[int, int])
extra (bool)

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:

x (Tensor)
dim (int | tuple[int, int])
extra (bool)

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:

x (Tensor)
dim (int | tuple[int, int])
extra (bool)

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:

x (Tensor)
dim (int | tuple[int, int])
extra (bool)

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:

x (Tensor)
dim (int | tuple[int, int])
extra (bool)

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:

x (Tensor)
dim (int | tuple[int, int])
extra (bool)

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:

x (Tensor)
dim (int | tuple[int, int])
extra (bool)

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:

x (Tensor)
dim (int | tuple[int, int])
extra (bool)

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:

x (Tensor)
dim (int | tuple[int, int])
extra (bool)

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:

x (Tensor)
dim (int | tuple[int, int])
extra (bool)

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:

x (Tensor)
dim (int | tuple[int, int])
extra (bool)

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:

device (device | None)
dtype (dtype | None)

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.

property nao: int#

property nat: int#

property nbatch: int | None#

property nsh: int#

orbital_index: Tensor#: Offset index for starting the next orbital block

property orbital_index_cart: Tensor#

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

property orbitals_per_shell_cart: Tensor#

property orbitals_to_atom: Tensor#

property orbitals_to_atom_cart: Tensor#

orbitals_to_shell: Tensor#: Mapping of orbitals to shells

property orbitals_to_shell_cart: Tensor#

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

dxtb.IndexHelper

Contents

dxtb.IndexHelper#