Integrals and GTO basis (low level)

The analytic integral matrices (Obara-Saika / McMurchie-Davidson, all jit/vmap/grad-friendly) and the GTO basis container they act on. Most users never call these directly.

dftax.integrals.overlap_matrix(basis: BasisData, plan: tuple | None = None) -> Float[Array, 'nao nao']

Compute the overlap matrix S_μν in the AO basis.

Pure JAX, fully differentiable w.r.t. basis.centers. Delegates to the shell-class-bucketed engine (see :mod:dftax.integrals.eri3c_bucketed); plan is a static :func:~dftax.integrals.eri3c_bucketed.plan_pairs skeleton, derived here when the basis metadata is concrete.

Parameters:

Name	Type	Description	Default
basis	BasisData	BasisData from extract_basis_data(mol).	required

Returns:

Type	Description
Float[Array, 'nao nao']	S matrix, shape (nao, nao) where nao = nao_cart (or nao_sph if cart2sph).

dftax.integrals.kinetic_matrix(basis: BasisData, plan: tuple | None = None) -> Float[Array, 'nao nao']

Compute the kinetic energy matrix T_μν in the AO basis.

Pure JAX, fully differentiable w.r.t. basis.centers. Delegates to the shell-class-bucketed engine (see :mod:dftax.integrals.eri3c_bucketed).

Parameters:

Name	Type	Description	Default
basis	BasisData	BasisData from extract_basis_data(mol).	required

Returns:

Type	Description
Float[Array, 'nao nao']	T matrix, shape (nao, nao).

dftax.integrals.nuclear_attraction_matrix(basis: BasisData, atom_coords: Float[Array, 'n_atoms 3'], atom_charges: Float[Array, n_atoms], plan: tuple | None = None) -> Float[Array, 'nao nao']

Compute nuclear attraction matrix V_μν in the AO basis.

V_μν = -Σ_A Z_A ∫ χ_μ(r) · 1/|r-R_A| · χ_ν(r) dr

Pure JAX, fully differentiable w.r.t. basis.centers and atom_coords. Delegates to the shell-class-bucketed engine (see :mod:dftax.integrals.eri3c_bucketed): the flat per-element build held every (pair, nucleus) Hermite table at molecule-padded sizes at once and owned the KS build's memory peak (27.5 GiB for ethanol/def2-svp).

Parameters:

Name	Type	Description	Default
basis	BasisData	BasisData from extract_basis_data(mol).	required
atom_coords	Float[Array, 'n_atoms 3']	Nuclear positions, shape (n_atoms, 3).	required
atom_charges	Float[Array, n_atoms]	Nuclear charges, shape (n_atoms,).	required
plan	tuple | None	static pair skeleton from :func:~dftax.integrals.eri3c_bucketed.plan_pairs; required when this build is traced with a fully-traced BasisData, derived here otherwise.	None

Returns:

Type	Description
Float[Array, 'nao nao']	V matrix, shape (nao, nao).

dftax.integrals.nuclear_repulsion

Nuclear repulsion energy V_nn = Σ_{A<B} Z_A Z_B / |R_A - R_B|.

Trivially differentiable w.r.t. nuclear coordinates via JAX autodiff.

__cached__ module-attribute

str(object='') -> str str(bytes_or_buffer[, encoding[, errors]]) -> str

Create a new string object from the given object. If encoding or errors is specified, then the object must expose a data buffer that will be decoded using the given encoding and error handler. Otherwise, returns the result of object.str() (if defined) or repr(object). encoding defaults to 'utf-8'. errors defaults to 'strict'.

__doc__ module-attribute

str(object='') -> str str(bytes_or_buffer[, encoding[, errors]]) -> str

Create a new string object from the given object. If encoding or errors is specified, then the object must expose a data buffer that will be decoded using the given encoding and error handler. Otherwise, returns the result of object.str() (if defined) or repr(object). encoding defaults to 'utf-8'. errors defaults to 'strict'.

__file__ module-attribute

str(object='') -> str str(bytes_or_buffer[, encoding[, errors]]) -> str

Create a new string object from the given object. If encoding or errors is specified, then the object must expose a data buffer that will be decoded using the given encoding and error handler. Otherwise, returns the result of object.str() (if defined) or repr(object). encoding defaults to 'utf-8'. errors defaults to 'strict'.

__name__ module-attribute

str(object='') -> str str(bytes_or_buffer[, encoding[, errors]]) -> str

Create a new string object from the given object. If encoding or errors is specified, then the object must expose a data buffer that will be decoded using the given encoding and error handler. Otherwise, returns the result of object.str() (if defined) or repr(object). encoding defaults to 'utf-8'. errors defaults to 'strict'.

__package__ module-attribute

str(object='') -> str str(bytes_or_buffer[, encoding[, errors]]) -> str

Create a new string object from the given object. If encoding or errors is specified, then the object must expose a data buffer that will be decoded using the given encoding and error handler. Otherwise, returns the result of object.str() (if defined) or repr(object). encoding defaults to 'utf-8'. errors defaults to 'strict'.

nuclear_repulsion(coords: Float[Array, 'n_atoms 3'], charges: Float[Array, n_atoms]) -> Shaped[Array, '']

Compute nuclear repulsion energy.

Parameters:

Name	Type	Description	Default
coords	Float[Array, 'n_atoms 3']	Nuclear positions in Bohr, shape (n_atoms, 3).	required
charges	Float[Array, n_atoms]	Nuclear charges, shape (n_atoms,).	required

Returns:

Type	Description
Shaped[Array, '']	Nuclear repulsion energy in Hartree (scalar).

dftax.integrals.eri2c_matrix(aux_basis: BasisData, omega: float | None = None, plan: tuple | None = None) -> Float[Array, 'n_aux n_aux']

Compute 2-center Coulomb matrix J_{PQ} = (P|Q).

Pure JAX, fully differentiable w.r.t. aux_basis.centers. Delegates to the shell-class-bucketed engine (see :mod:dftax.integrals.eri3c_bucketed).

Parameters:

Name	Type	Description	Default
aux_basis	BasisData	BasisData for auxiliary basis (from extract_basis_data(auxmol)).	required
omega	float | None	None for the Coulomb kernel, a float for the long-range one.	None
plan	tuple | None	static pair skeleton from :func:~dftax.integrals.eri3c_bucketed.plan_pairs; required when traced with a fully-traced BasisData, derived here otherwise.	None

Returns:

Type	Description
Float[Array, 'n_aux n_aux']	J matrix, shape (n_aux, n_aux) in spherical harmonics.

dftax.integrals.eri3c_matrix(basis: BasisData, aux_basis: BasisData, omega: float | None = None, plan: tuple | None = None) -> Float[Array, 'nao nao n_aux']

Compute 3-center ERI tensor (μν|P) in the AO basis.

Pure JAX, fully differentiable w.r.t. basis.centers and aux_basis.centers. Delegates to the shell-class-bucketed engine (:func:~dftax.integrals.eri3c_bucketed.eri3c_matrix_bucketed): one right-sized kernel per (l_a, l_b, l_aux) class instead of a molecule-padded per-element build (50x on ethanol/def2-svp, sub-GiB peak, identical to machine precision).

Parameters:

Name	Type	Description	Default
basis	BasisData	BasisData for primary AO basis (from extract_basis_data(mol)).	required
aux_basis	BasisData	BasisData for auxiliary basis (from extract_basis_data(auxmol)).	required
omega	float | None	None for the Coulomb kernel, a float for the long-range erf(ω·r₁₂)/r₁₂ kernel (range-separated hybrids).	None
plan	tuple | None	static bucket skeleton from :func:~dftax.integrals.eri3c_bucketed.plan_eri3c; required when this build is traced with a fully-traced BasisData (the jitted _build_integrals passes it), derived here otherwise.	None

Returns:

Type	Description
Float[Array, 'nao nao n_aux']	(μν|P) tensor, shape (nao, nao, n_aux) in spherical harmonics.

dftax.integrals.multipole.dipole_matrices(basis: BasisData, origin=(0.0, 0.0, 0.0)) -> Float[Array, '3 nao nao']

Dipole integral matrices ⟨μ|(r-origin)|ν⟩ for the x, y, z components.

Differentiable w.r.t. basis.centers and origin. Returns shape (3, nao, nao) in the same AO ordering as :func:~dftax.integrals.overlap_matrix (spherical if basis.cart2sph is set).

dftax.energy.gto.BasisData

Precomputed GTO basis data extracted from a PySCF mol object.

All arrays are JAX arrays; construction is done once on CPU via extract_basis_data(mol).

Attributes:

Name	Type	Description
centers		Atom centre for each AO, shape (nao_cart, 3).
exponents		Primitive exponents, zero-padded, shape (nao_cart, max_prim).
coefficients		Final contraction coefficients (primitive-normed * contracted-normed), zero-padded, shape (nao_cart, max_prim).
angular		Cartesian angular momenta (lx, ly, lz), shape (nao_cart, 3).
cart2sph		Cartesian-to-spherical transformation, shape (nao_cart, nao_sph), or None when the molecule uses Cartesian GTOs.
max_l		Maximum total angular momentum in the basis (static int). Lets the integral builders size their recursion to the molecule instead of the global g-type cap.

__abstractclassvars__ class-attribute

Build an immutable unordered collection of unique elements.

__abstractmethods__ class-attribute

Build an immutable unordered collection of unique elements.

__abstractvars__ class-attribute

Build an immutable unordered collection of unique elements.

__annotations__ class-attribute

dict() -> new empty dictionary dict(mapping) -> new dictionary initialized from a mapping object's (key, value) pairs dict(iterable) -> new dictionary initialized as if via: d = {} for k, v in iterable: d[k] = v dict(**kwargs) -> new dictionary initialized with the name=value pairs in the keyword argument list. For example: dict(one=1, two=2)

__dataclass_fields__ class-attribute

dict() -> new empty dictionary dict(mapping) -> new dictionary initialized from a mapping object's (key, value) pairs dict(iterable) -> new dictionary initialized as if via: d = {} for k, v in iterable: d[k] = v dict(**kwargs) -> new dictionary initialized with the name=value pairs in the keyword argument list. For example: dict(one=1, two=2)

__dataclass_params__ class-attribute

__doc__ class-attribute

str(object='') -> str str(bytes_or_buffer[, encoding[, errors]]) -> str

Create a new string object from the given object. If encoding or errors is specified, then the object must expose a data buffer that will be decoded using the given encoding and error handler. Otherwise, returns the result of object.str() (if defined) or repr(object). encoding defaults to 'utf-8'. errors defaults to 'strict'.

__firstlineno__ class-attribute

int([x]) -> integer int(x, base=10) -> integer

Convert a number or string to an integer, or return 0 if no arguments are given. If x is a number, return x.int(). For floating-point numbers, this truncates towards zero.

If x is not a number or if base is given, then x must be a string, bytes, or bytearray instance representing an integer literal in the given base. The literal can be preceded by '+' or '-' and be surrounded by whitespace. The base defaults to 10. Valid bases are 0 and 2-36. Base 0 means to interpret the base from the string as an integer iteral.

int('0b100', base=0) 4

__match_args__ class-attribute

Built-in immutable sequence.

If no argument is given, the constructor returns an empty tuple. If iterable is specified the tuple is initialized from iterable's items.

If the argument is a tuple, the return value is the same object.

__module__ class-attribute

str(object='') -> str str(bytes_or_buffer[, encoding[, errors]]) -> str

Create a new string object from the given object. If encoding or errors is specified, then the object must expose a data buffer that will be decoded using the given encoding and error handler. Otherwise, returns the result of object.str() (if defined) or repr(object). encoding defaults to 'utf-8'. errors defaults to 'strict'.

__static_attributes__ class-attribute

Built-in immutable sequence.

If no argument is given, the constructor returns an empty tuple. If iterable is specified the tuple is initialized from iterable's items.

If the argument is a tuple, the return value is the same object.

max_l class-attribute

int([x]) -> integer int(x, base=10) -> integer

Convert a number or string to an integer, or return 0 if no arguments are given. If x is a number, return x.int(). For floating-point numbers, this truncates towards zero.

If x is not a number or if base is given, then x must be a string, bytes, or bytearray instance representing an integer literal in the given base. The literal can be preceded by '+' or '-' and be surrounded by whitespace. The base defaults to 10. Valid bases are 0 and 2-36. Base 0 means to interpret the base from the string as an integer iteral.

int('0b100', base=0) 4

__init__(centers: Float[Array, 'nao_cart 3'], exponents: Float[Array, 'nao_cart max_prim'], coefficients: Float[Array, 'nao_cart max_prim'], angular: Int[Array, 'nao_cart 3'], cart2sph: Float[Array, 'nao_cart nao_sph'] | None, max_l: int = 4) -> None

Initialize self. See help(type(self)) for accurate signature.

dftax.energy.gto.extract_basis_data(mol: object) -> BasisData

Extract GTO basis data from a PySCF Mole object.

This is a CPU-side function called once at setup. It returns a BasisData whose arrays live on the JAX default device.

Parameters:

Name	Type	Description	Default
mol	object	A built PySCF gto.Mole object.	required

Returns:

Type	Description
BasisData	BasisData suitable for eval_gto.

dftax.energy.gto.eval_gto(basis: BasisData, r: Float[Array, 3]) -> Float[Array, nao]

Evaluate contracted GTO basis functions at point r.

Pure JAX, no PySCF calls, no pure_callback. Fully differentiable via JAX autodiff and compatible with jit and vmap.

Parameters:

Name	Type	Description	Default
basis	BasisData	Precomputed basis data from extract_basis_data(mol).	required
r	Float[Array, 3]	3D coordinate in Bohr, shape (3,).	required

Returns:

Type	Description
Float[Array, nao]	AO values, shape (nao,), matching dft.numint.eval_ao(mol, r[None])[0].