Solvers and results

Five verbs over one built functional, one result type: scf (DIIS, with optional accel=adiis() far-from-convergence extrapolation and smearing=fermi() fractional occupations reporting the Mermin free energy), newton (trust-region second order on orbital rotations, saddle-robust Steihaug-Toint steps), roks (its shared-orbital restricted open-shell variant), minimize (direct orbital optimization with any optax optimizer), and gdm (direct minimization by L-BFGS over occupied-virtual rotations, evaluated through the occupied orbitals). implicit_density exposes the converged density as a differentiable function of the functional (CPHF via custom_vjp).

All solvers (and scf_batched) take a guess=: an initial-density spec from core / sad / minao / sap, or an explicit (nspin, nao, nao) density array for warm restarts. The guess changes the iteration count, never the converged fixed point. The default is minao(); the core Hamiltonian is the fallback for a raw System (no element identities) or an element the minimal basis does not reach.

dftax.ks.scf.scf(ks: KS, *, max_iter: int = 128, e_tol: float = 1e-08, d_tol: float = 1e-06, diis_space: int = 8, lindep_thresh: float = 1e-07, level_shift: float = 0.0, guess: GuessSpec | Array | None = None, accel: ADIISSpec | None = None, smearing: FermiSpec | None = None, verbose: bool = False) -> KSResult

Run KS SCF to self-consistency (restricted and spin-polarized alike).

Parameters:

Name	Type	Description	Default
ks	KS	the built :class:~dftax.ks.energy.KS energy functional.	required
max_iter	int	maximum SCF iterations.	128
e_tol	float	convergence threshold on the total-energy change (Ha). An energy evaluation has a floating-point floor of roughly 5e-11·|E|, essentially all of it in the density-fitted Coulomb term (the RI metric pseudo-inverse retains directions down to 1e-7 of the largest eigenvalue, so contracting against the explicit inverse is that ill-conditioned by construction; the XC quadrature and the one-electron terms are clean to 1e-14). Measured by sweeping the density along a line and fitting the parabola the energy is analytically supposed to be: 1.06e-8 Ha on cubane, 2.30e-8 on bicyclo[2.2.2]octane, 2.07e-8 on adamantane, 3.52e-8 on coronene. A threshold below that is not a test but a coin flip, which is worth knowing before tightening this: on coronene, e_tol=1e-9 costs 102 iterations against 10 for 1e-8, and buys 8.4e-9 Ha. The floor is far below chemical accuracy and does not affect any energy difference; it only limits what this argument can usefully be set to, and it grows with the molecule while the default does not.	1e-08
d_tol	float	convergence threshold on the DIIS commutator norm. This is the orbital gradient, and unlike e_tol it has no such floor, so it is the criterion to tighten and the one that keeps working as the system grows. PySCF's convention for the pair is d_tol = sqrt(e_tol).	1e-06
diis_space	int	DIIS history depth (fixed buffer size).	8
lindep_thresh	float	overlap-eigenvalue cutoff for canonical orthonormalization.	1e-07
level_shift	float	Saunders-Hillier virtual level shift (Ha).	0.0
guess	GuessSpec | Array | None	initial density, a spec from :func:~dftax.ks.guess.core / :func:~dftax.ks.guess.sad / :func:~dftax.ks.guess.minao / :func:~dftax.ks.guess.sap, or an explicit (nspin, nao, nao) density array (warm restart). None is :func:~dftax.ks.guess.minao, falling back to the core Hamiltonian where the atomic densities cannot be built.	None
accel	ADIISSpec | None	an :func:adiis spec runs the far-from-convergence iterations with ADIIS (energy-model extrapolation, no oscillation), blending into Pulay DIIS near the fixed point; None is plain Pulay DIIS. When given, its space supersedes diis_space.	None
smearing	FermiSpec | None	a :func:fermi spec replaces the integer aufbau fill with Fermi-Dirac fractional occupations (small-gap systems, smooth energies through level crossings); None is the aufbau fill.	None
verbose	bool	print per-iteration energy / error (via jax.debug.print).	False

Example

ks = KS(mol, PBE())
res = scf(ks, e_tol=1e-9)
res = scf(ks, guess=sad())               # fewer iterations
res.e_tot, res.converged, res.P[0]       # P is spin-stacked

dftax.ks.scf.adiis(*, space: int = 8, switch: float = 0.1) -> ADIISSpec

ADIIS-accelerated SCF (Hu-Yang energy-model extrapolation).

Plain Pulay DIIS extrapolates on the commutator error and can oscillate or limit-cycle when the starting density is poor (transition metals, radicals, stretched geometries). ADIIS instead minimizes an interpolated energy model over the density/Fock history, which cannot overshoot, and switches to Pulay DIIS whenever the commutator norm is below switch (Pulay converges faster near the fixed point; the switch is re-entrant, so an error that grows again hands back to ADIIS).

Parameters:

Name	Type	Description	Default
space	int	history depth (shared by the ADIIS and Pulay buffers).	8
switch	float	commutator-norm threshold between the ADIIS and Pulay extrapolations.	0.1

Example

res = scf(ks, accel=adiis())             # robust far-from-convergence
res = scf(ks, accel=adiis(switch=1e-2))  # trust ADIIS longer

dftax.ks.scf.fermi(*, sigma: float = 0.01) -> FermiSpec

Fermi-Dirac smearing of the SCF occupations.

Occupations become f_i = w / (1 + exp((eps_i - mu) / sigma)) with the chemical potential mu solved per spin channel so the electron count is conserved. Small-gap and metallic-like systems (transition metals, stretched bonds, clusters) converge where integer aufbau occupations flip-flop between quasi-degenerate configurations, and the occupations become a smooth function of the orbital energies, which keeps the energy differentiable through level crossings.

The reported e_tot is the Mermin free energy A = E - sigma·S (the Kohn-Sham energy at the smeared density minus the electronic-entropy term), which is the quantity the finite-temperature SCF makes stationary and whose gradient is the force. The entropy term is returned as KSResult.ts so the KS energy (e_tot + ts) and the sigma -> 0 estimate (e_tot + 0.5·ts) are recoverable; taking sigma to zero sends the free energy to the integer-occupation ground-state energy.

Parameters:

Name	Type	Description	Default
sigma	float	smearing width in Hartree (0.01 Ha ~ 3160 K electronic temperature; typical range 0.001 to 0.02).	0.01

Example

res = scf(ks, smearing=fermi(sigma=0.005))

dftax.ks.minimize.minimize(ks: KS, optimizer: optax.GradientTransformation | None = None, *, max_steps: int = 2000, g_tol: float = 1e-06, check_every: int = 10, Z0: Array | tuple | None = None, guess: GuessSpec | Array | None = None, verbose: bool = False) -> KSResult

Minimize the KS energy directly over orthonormalized coefficients.

Parameters:

Name	Type	Description	Default
ks	KS	the built :class:~dftax.ks.energy.KS energy functional.	required
optimizer	GradientTransformation | None	any optax.GradientTransformation; default optax.adam(0.3).	None
max_steps	int	optimizer step budget.	2000
g_tol	float	stop when the global gradient norm falls below this.	1e-06
check_every	int	host-side convergence checks happen every this many steps (each check blocks on the device), so the reported n_iter can overshoot the true crossing by up to check_every - 1 steps. verbose checks every step.	10
Z0	Array | tuple | None	optional initial coefficient guess, one (nao, nocc_σ) array per channel (a bare array is accepted for a closed shell); the default follows guess below.	None
guess	GuessSpec | Array | None	alternative to Z0: a spec from :func:~dftax.ks.guess.core / :func:~dftax.ks.guess.sad / :func:~dftax.ks.guess.minao / :func:~dftax.ks.guess.sap or an explicit (nspin, nao, nao) density; the starting orbitals are the aufbau eigenvectors of the Fock built at that density.	None
verbose	bool	print per-step energy and gradient norm.	False

Example

res = minimize(ks)                             # optax.adam(0.3)
res = minimize(ks, optax.chain(
    optax.clip_by_global_norm(1.0), optax.adam(0.1)), g_tol=1e-7)

dftax.ks.gdm.gdm(ks: KS, *, max_iter: int = 200, e_tol: float = 1e-08, d_tol: float = 1e-06, history: int = 20, max_step: float = 0.5, gap_floor: float = 0.05, lindep_thresh: float = 1e-07, guess: GuessSpec | jax.Array | None = None, orbital_xc: bool = True, verbose: bool = False) -> KSResult

Minimize the KS energy over orbital rotations (restricted and spin-polarized alike).

Parameters:

Name	Type	Description	Default
ks	KS	the built :class:~dftax.ks.energy.KS energy functional.	required
max_iter	int	budget of energy evaluations (each is one Fock build).	200
e_tol	float	convergence threshold on the energy change of an accepted step (Ha); same meaning and floor as in :func:~dftax.ks.scf.scf.	1e-08
d_tol	float	convergence threshold on the commutator norm, measured exactly as :func:~dftax.ks.scf.scf measures it, so the two solvers stop at the same point.	1e-06
history	int	L-BFGS memory depth.	20
max_step	float	largest rotation angle (rad) any one step may take.	0.5
gap_floor	float	smallest orbital-energy gap (Ha) the preconditioner uses, which keeps it positive where the current orbitals are not in aufbau order.	0.05
lindep_thresh	float	overlap-eigenvalue cutoff for the orthonormalizer.	1e-07
guess	GuessSpec | Array | None	as for :func:~dftax.ks.scf.scf.	None
orbital_xc	bool	evaluate through the occupied orbitals (:meth:~dftax.ks.energy.KS.energy_and_grad_orbitals), building the Fock only at the guess and at reference resets; False builds it at every evaluation.	True
verbose	bool	print every energy evaluation.	False

n_iter of the result counts energy evaluations, the unit that costs.

dftax.ks.newton.newton(ks: KS, *, max_iter: int = 64, g_tol: float = 1e-06, e_tol: float = 1e-10, cg_iters: int = 32, trust: float = 0.4, guess: GuessSpec | Array | None = None, verbose: bool = False) -> KSResult

Second-order SCF (trust-region Newton on orbital rotations).

Converges quadratically near a minimum and handles densities that DIIS and ADIIS limit-cycle on (small-gap transition metals, stretched bonds). One Newton iteration costs a CG solve of Hessian-vector products, each the price of a couple of Fock builds, so for benign systems :func:~dftax.ks.scf.scf is faster; reach for this when robustness is the constraint, typically warm-started from a stalled SCF density (newton(ks, guess=res.P)).

Parameters:

Name	Type	Description	Default
ks	KS	the built :class:~dftax.ks.energy.KS energy functional.	required
max_iter	int	maximum Newton iterations.	64
g_tol	float	convergence threshold on the orbital-gradient norm.	1e-06
e_tol	float	convergence threshold on the energy change (Ha).	1e-10
cg_iters	int	CG iterations for the Newton step (Hessian-vector solves).	32
trust	float	initial trust radius on the rotation norm (radians); halved on an uphill proposal, doubled (capped at 1) on success.	0.4
guess	GuessSpec | Array | None	initial density, same forms as :func:~dftax.ks.scf.scf.	None
verbose	bool	print per-iteration energy / gradient norm.	False

Example

res = scf(ks, max_iter=200)              # limit-cycles on a hard case
res = newton(ks, guess=res.P)            # quadratic cleanup

dftax.ks.newton.roks(ks: KS, *, max_iter: int = 64, g_tol: float = 1e-06, e_tol: float = 1e-10, cg_iters: int = 32, trust: float = 0.4, guess: GuessSpec | Array | None = None, verbose: bool = False) -> KSResult

Restricted open-shell KS via shared-orbital Newton.

One set of spatial orbitals for both spin channels (alpha fills the first nocc[0], beta the first nocc[1]), so the wavefunction is a spin eigenfunction with none of UKS's spin contamination. Implemented as :func:newton over rotations of the single shared orbital set: kappa is masked to the doubly-occupied/singly-occupied/virtual inter-block rotations that change the density, and the constraint that the beta space sits inside the alpha space holds by construction. Requires a spin-polarized ks (spin != 0).

Parameters:

Name	Type	Description	Default
ks	KS	the built spin-polarized :class:~dftax.ks.energy.KS.	required
max_iter	int	maximum Newton iterations.	64
g_tol	float	convergence threshold on the orbital-gradient norm.	1e-06
e_tol	float	convergence threshold on the energy change (Ha).	1e-10
cg_iters	int	CG iterations for the Newton step.	32
trust	float	initial trust radius on the rotation norm.	0.4
guess	GuessSpec | Array | None	initial density, same forms as :func:~dftax.ks.scf.scf.	None
verbose	bool	print per-iteration energy / gradient norm.	False

Example

ks = KS(mol, PBE(), spin=2)              # triplet
res = roks(ks)                           # E >= UKS, no contamination

dftax.ks.scf.KSResult

Outcome of a KS solve (SCF or direct minimization), spin-stacked.

All orbital/density fields carry a leading nspin axis (nspin = len(nocc)): a closed shell is nspin = 1 (P[0] doubly occupied), a spin-polarized system is nspin = 2 (α, β). Batched solves return the distinct :class:~dftax.ks.batched.BatchedResult instead.

__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

ts class-attribute

Convert a string or number to a floating-point number, if possible.

__eq__(other)

__init__(e_tot: float, e_elec: float, converged: bool, n_iter: int, nocc: tuple[int, ...], mo_energy: Float[Array, 'nspin nmo'], mo_coeff: Float[Array, 'nspin nao nmo'], P: Float[Array, 'nspin nao nao'], ts: float = 0.0) -> None

__repr__()

dftax.ks.implicit.implicit_density(ks: KS)

Converged closed-shell density P* for the assembled functional ks, differentiable w.r.t. ks (and hence anything it was assembled from) by implicit differentiation of the SCF fixed point. Restricted (closed-shell) only.

Initial guesses

dftax.ks.guess.core() -> CoreSpec

Core-Hamiltonian guess: aufbau fill of the hcore eigenvectors.

The dftax default (guess=None); ignores electron repulsion, so it is the weakest standard guess: fine for small closed-shell systems, slow to converge (or saddle-prone under :func:~dftax.ks.minimize.minimize) for larger ones.

dftax.ks.guess.sad() -> SADSpec

Superposition of atomic densities from solved atoms.

Per unique element, a spherically-averaged (fractional-occupation) atomic Hartree-Fock is converged in the molecule's own basis, and the atomic densities are assembled block-diagonally. The strongest density-based guess here; the per-element solves are tiny and cached.

dftax.ks.guess.minao(basis: str = 'sto-3g') -> MinAOSpec

Superposition of atomic densities projected from a minimal basis.

Per unique element, tabulated ground-state occupations are placed on the minimal-basis AOs and projected onto the computational basis through the cross overlap. Cheaper than :func:sad (no atomic solves), nearly as good.

Parameters:

Name	Type	Description	Default
basis	str	minimal basis holding the atomic ground states, any Basis Set Exchange name with one shell per occupied (n, l) ("sto-3g", "mini", "ano-rcc-mb").	'sto-3g'

dftax.ks.guess.sap(fit: str = 'sap_helfem_large', *, grid: Becke | None = None) -> SAPSpec

Superposition of atomic potentials: F0 = T + V_SAP.

A potential-based guess (no density superposition): the screened effective potential of each atom, Gaussian-fitted, summed over atoms and integrated on a Becke grid. Robust and cheap; the only guess here whose quality does not depend on a converged atomic density.

Parameters:

Name	Type	Description	Default
fit	str	the Gaussian potential fit, a Basis Set Exchange name ("sap_helfem_large", "sap_helfem_small", "sap_grasp_large", "sap_grasp_small").	'sap_helfem_large'
grid	Becke | None	quadrature for the potential matrix, a :func:~dftax.grid.becke spec; None uses the default grid.	None