Building a calculation

One builder captures every choice as a value: the system, the functional, and the grid/Coulomb/spin/mesh/dispersion specs. See All of dftax for the tour.

dftax.ks.energy.KS

Spin-stacked KS total energy as a differentiable function of P.

Built directly from a system + functional + backend/grid values (see the module docstring). P has shape (nspin, nao, nao) with nspin = len(self.nocc): one doubly-occupied channel (P[0] = 2ΣCCᵀ) for a closed shell, two unit-occupation channels (P[σ] = ΣC_σC_σᵀ) for a spin-polarized system.

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

__init__(system: Molecule | System | object, xc: XCFunctional, *, grid: Becke | Points | tuple | None = None, coulomb: ExactSpec | DFSpec | None = None, spin: int | None = None, mesh: MeshSpec | None = None, dispersion: D3BJSpec | None = None)

Build the energy functional.

Parameters:

Name	Type	Description	Default
system	Molecule | System | object	a :class:~dftax.system.molecule.Molecule, a PySCF Mole (setup only), or a raw :class:System.	required
xc	XCFunctional	the exchange-correlation functional (e.g. PBE()).	required
grid	Becke | Points | tuple | None	quadrature choice, a :func:~dftax.grid.becke spec (default), an explicit (coords, weights) tuple, or a :func:~dftax.grid.points spec (explicit grid + streaming chunk).	None
coulomb	ExactSpec | DFSpec | None	Coulomb/exchange backend, :func:~dftax.ks.terms.df (default: def2-universal-jkfit, auto materialize/stream) or :func:~dftax.ks.terms.exact (O(N⁴); also the raw-System fallback).	None
spin	int | None	None infers from the system (closed shell → restricted); an explicit spin (= 2S, including 0) requests spin-polarized α/β channels.	None
mesh	MeshSpec | None	a :func:~dftax.ks.shard.mesh spec to shard the calculation across a device mesh (the XC quadrature and the DF 3-center tensor; the dense nao² matrices stay replicated). None = single device.	None
dispersion	D3BJSpec | None	a :func:~dftax.energy.d3.d3bj spec adds the two-body D3(BJ) correction to total() (P-independent, so the SCF is untouched; forces get its gradient by autodiff). Parameters resolve from the functional's name unless given explicitly.	None

Example

ks = KS(mol, PBE0(),
        grid=becke(75, 302),
        coulomb=df("def2-universal-jkfit"))

e_xc(P: Float[Array, 'nspin nao nao']) -> Scalar

Exchange-correlation energy ∫ ε_xc ρ (DFT part only).

electronic(P: Float[Array, 'nspin nao nao']) -> Scalar

Electronic energy Tr(P_tot·Hcore) + E_J + a_x·E_x^exact + E_xc.

total(P: Float[Array, 'nspin nao nao']) -> Scalar

Total KS energy (electronic + nuclear repulsion + dispersion).

energy_and_fock(P: Float[Array, 'nspin nao nao'], idempotent: bool = True) -> tuple[Scalar, Float[Array, 'nspin nao nao']]

(E_total, F) for one density, sharing the work between them.

Computing the two separately walks the quadrature twice, because the streaming backends checkpoint their grid blocks and grad rematerializes them. Each term is therefore asked for both at once (:meth:~dftax.ks.terms.XCTerm.energy_and_potential, :meth:~dftax.ks.terms.CoulombTerm.energy_and_potential).

idempotent says whether P is an integer-occupation projector: true of an aufbau density, false under Fermi smearing. The Coulomb term falls back to reverse mode when told otherwise.

The result is the same (total(P), sym(∂E/∂P)) as the two separate calls, which total / electronic still provide for every other consumer (forces, the Hessian, Newton, direct minimization).

energy_and_grad_orbitals(Cs, idempotent: bool = True)

(E_total, ∂E/∂C_σ) for the occupied orbitals of each channel.

The density is P_σ = C_σ diag(w_σ) C_σᵀ with the occupations of :meth:occupations (2 closed shell, 1 per spin channel), so ∂E/∂C_σ = 2 F_σ C_σ diag(w_σ) with F = sym(∂E/∂P). The one-electron and two-electron parts take that route, through their (cheap) potentials; the XC part, almost all of an iteration, comes straight from the orbitals (:meth:~dftax.ks.terms.XCTerm.energy_and_grad_orbitals), which never forms the nao x nao XC potential. What a direct minimizer (:func:~dftax.ks.gdm.gdm) needs, and nothing more.

occupations()

Occupation of each channel's occupied orbitals: 2 each for a closed shell, 1 each per spin channel. A padded system (:mod:dftax.ks.batched_gdm) gives its padding slots 0.

density(P: Float[Array, 'nspin nao nao']) -> tuple[Float[Array, ng], Float[Array, 'ng 3']]

Total electron density and its gradient on the grid from P.

dftax.energy.d3.d3bj(method: str | None = None, *, atm: bool = False) -> D3BJSpec

D3(BJ) dispersion correction (two-body, optionally with ATM).

Parameters:

Name	Type	Description	Default
method	str | None	damping-parameter set name ("pbe", "pbe0", "b3lyp", "cam-b3lyp", "r2scan"); None resolves it from the functional the KS is built with.	None
atm	bool	add the Axilrod-Teller-Muto three-body term (s9 = 1); off by default, matching the common two-body D3(BJ) convention. The term is repulsive, O(nat³), and independent of the BJ damping parameters (it uses the pairwise van-der-Waals radii).	False

dftax.energy.d4.d4(method: str | None = None, *, atm: bool = False) -> D4Spec

Charge-dependent D4 dispersion correction.

Parameters:

Name	Type	Description	Default
method	str | None	damping-parameter set name ("pbe", "pbe0", "b3lyp", "cam-b3lyp", "r2scan"); None resolves it from the functional the KS is built with.	None
atm	bool	add the Axilrod-Teller-Muto three-body term on the D4 coefficients (s9 = 1); off by default.	False

dftax.system.molecule.Molecule

A molecular system: atoms, geometry, charge/spin, and a basis name.

__dict__ class-attribute

Read-only proxy of a mapping.

__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

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

__weakref__ property

list of weak references to the object

nelectron property

Total electron count (sum of Z minus the total charge).

__init__(symbols: list[str], coords_bohr: np.ndarray, basis: str, charge: int = 0, spin: int = 0, spherical: bool = False)

from_xyz(atom: str, basis: str, *, unit: str = 'angstrom', charge: int = 0, spin: int = 0, spherical: bool = False) -> Molecule classmethod

Build from a PySCF-style atom string (Angstrom by default).

atom_coords() -> np.ndarray

Nuclear coordinates in Bohr, shape (n_atoms, 3).

atom_charges() -> np.ndarray

Nuclear charges (atomic numbers), shape (n_atoms,).

dftax.ks.energy.System

A raw, already-built system: basis + geometry + electron/spin counts.

The low-level :class:KS input for callers that construct the basis themselves, e.g. forces, which rebuild the energy as a function of traced nuclear coordinates via a tree_at-recentered basis template. spin (= 2S) is the system's default spin; the :class:KS builder's spin argument overrides it.

__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

__dict__ class-attribute

Read-only proxy of a mapping.

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

__weakref__ property

list of weak references to the object

spin 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

__delattr__(name)

__eq__(other)

__hash__()

__init__(basis: BasisData, coords: Array, charges: Array, nelec: int, spin: int = 0) -> None

__repr__()

__setattr__(name, value)