API reference
Drivers
Unified entry point: dispatch to RKS or UKS by spin.
spin (= 2S) defaults to the system's spin (mol.spin for a PySCF
Mole, Molecule.spin for the native type). By default a closed shell
(spin 0) runs restricted and anything else runs unrestricted; pass
restricted=True/False to force one path. Remaining keyword arguments are
forwarded to the chosen runner.
Run a restricted (closed-shell) RKS-DFT calculation.
Parameters:
| Name | Type | Description | Default |
|---|---|---|---|
system
|
a native :class: |
required | |
xc
|
XCFunctional
|
an :class: |
required |
auxbasis
|
str | None
|
optional density-fitting auxiliary basis name (e.g.
|
None
|
grid
|
tuple | None
|
optional |
None
|
grid_chunk
|
int | None
|
if set, stream the XC grid in chunks of this many points (O(chunk·nao) grid memory) instead of materializing the AO grid. |
None
|
df_chunk
|
int | None
|
if set (with |
None
|
eri_screen
|
float | None
|
if set (exact path, no |
None
|
n_radial, lebedev
|
native Becke-grid quality (when |
required | |
grid_level
|
int
|
PySCF grid level (when |
3
|
**scf_kwargs
|
forwarded to :func: |
{}
|
Run an unrestricted (open-shell) UKS-DFT calculation.
Mirrors :func:run_rks; spin (= 2S = nα − nβ) defaults to the system's
own spin. Backends: exact ERI, materialized DF, or streamed DF (df_chunk
+ df_screen: RI-J + per-spin RI-K) and streamed XC grid (grid_chunk).
**scf_kwargs go to :func:~dftax.ks.scf_uks.uks_scf.
Energy functionals
Bases: Module
Closed-shell KS total energy as a differentiable function of P.
from_pyscf(mol, xc, grid_coords, grid_weights, auxbasis=None, grid_chunk=None, df_chunk=None, eri_screen=None, exact_stream=False, df_screen=None)
classmethod
Build from a PySCF Mole (setup only) and a quadrature grid.
PySCF is used here solely to parse the orbital basis and supply nuclear
geometry/charges; nothing PySCF enters the compute path. auxbasis
(a basis-set name) enables density fitting; grid_chunk streams the XC
grid and df_chunk streams RI-J (both memory-light for large systems).
eri_screen (exact path) sets a Cauchy-Schwarz threshold to skip
negligible ERI quartets.
from_molecule(mol, xc, grid_coords, grid_weights, auxbasis=None, spherical=False, grid_chunk=None, df_chunk=None, eri_screen=None, exact_stream=False, df_screen=None)
classmethod
Build from a native :class:~dftax.system.molecule.Molecule (no PySCF).
auxbasis (a basis-set name, e.g. "def2-universal-jkfit") enables
density fitting. spherical=True uses spherical-harmonic orbitals
(standard for cc-pVXZ/def2; required to match a spherical reference for
l>=2 bases). grid_chunk streams the XC grid and df_chunk streams
RI-J (both memory-light). eri_screen (exact path) sets a
Cauchy-Schwarz threshold to skip negligible ERI quartets.
density(P)
Electron density and its gradient on the grid from P.
e_xc(P)
Exchange-correlation energy ∫ ε_xc ρ (DFT part only).
Uses precomputed AO grid values when materialized; otherwise streams the grid in chunks (recomputing AO per chunk) for O(chunk·nao) memory.
electronic(P)
Electronic energy Tr(P·Hcore) + E_J + a_x·E_x^exact + E_xc.
total(P)
Total KS energy (electronic + nuclear repulsion).
Bases: Module
Open-shell KS total energy as a differentiable function of (Pα, Pβ).
from_pyscf(mol, xc, grid_coords, grid_weights, auxbasis=None, spin=None, eri_screen=None, df_chunk=None, df_screen=None, grid_chunk=None)
classmethod
Build from a PySCF Mole (setup only) and a quadrature grid.
spin (= 2S) defaults to mol.spin. PySCF is used solely to parse
the basis and supply geometry/charges; nothing PySCF enters the compute
path. auxbasis enables density fitting; eri_screen (exact path)
sets a Cauchy-Schwarz screening threshold.
from_molecule(mol, xc, grid_coords, grid_weights, auxbasis=None, spherical=False, spin=None, eri_screen=None, df_chunk=None, df_screen=None, grid_chunk=None)
classmethod
Build from a native :class:~dftax.system.molecule.Molecule (no PySCF).
spin (= 2S) defaults to mol.spin. spherical=True uses
spherical-harmonic orbitals; auxbasis enables density fitting;
eri_screen (exact path) sets a Cauchy-Schwarz screening threshold.
e_xc(Pa, Pb)
Spin-polarized XC energy ∫ ε_xc(ρα,ρβ,∇ρα,∇ρβ) ρ_tot.
Mirrors the nan-safe double-where of :func:dftax.energy.grid.xc_energy
but stacks the two spin channels onto the last axis so the spin-aware
functionals receive (ρα,ρβ) (and (∇ρα,∇ρβ) for GGA) per point.
With grid_chunk set the AO grid is not materialized; the XC integral is
streamed instead (O(chunk·nao) memory) via :func:_streamed_e_xc_uks.
electronic(Pa, Pb)
Electronic energy Tr(P_tot·Hcore) + E_J + a_x·E_x^exact + E_xc.
total(Pa, Pb)
Total KS energy (electronic + nuclear repulsion).
Solvers, forces, minimization
Run closed-shell RKS SCF to self-consistency.
Parameters:
| Name | Type | Description | Default |
|---|---|---|---|
ks
|
RKS
|
the precomputed :class: |
required |
max_iter
|
int
|
maximum SCF iterations. |
128
|
e_tol
|
float
|
convergence threshold on the total-energy change (Ha). |
1e-08
|
d_tol
|
float
|
convergence threshold on the DIIS commutator norm. |
1e-06
|
diis_space
|
int
|
DIIS history depth (fixed buffer size). |
8
|
lindep_thresh
|
float
|
overlap-eigenvalue cutoff for canonical orthonormalization. |
1e-07
|
verbose
|
bool
|
print per-iteration energy / error (via jax.debug.print). |
False
|
Run open-shell UKS SCF to self-consistency.
Args mirror :func:dftax.ks.scf.rks_scf; the result carries per-spin MO
energies, MO coefficients, and density matrices as (α, β) tuples.
Nuclear forces F = -dE/dR (Ha/Bohr), shape (n_atom, 3).
Parameters:
| Name | Type | Description | Default |
|---|---|---|---|
mol
|
a native :class: |
required | |
xc
|
XCFunctional
|
the exchange-correlation functional. |
required |
C_occ
|
Float[Array, 'nao nocc']
|
converged occupied MO coefficients |
required |
auxbasis
|
str | None
|
optional density-fitting auxiliary basis (forces are then for the density-fitted energy surface). |
None
|
n_radial, lebedev
|
Becke-grid quality (match the energy calculation). |
required |
Nuclear forces F = -dE/dR (Ha/Bohr), shape (n_atom, 3).
Parameters:
| Name | Type | Description | Default |
|---|---|---|---|
mol
|
a native :class: |
required | |
xc
|
XCFunctional
|
the exchange-correlation functional. |
required |
Ca_occ, Cb_occ
|
converged occupied α/β MO coefficients (e.g.
|
required | |
auxbasis
|
str | None
|
optional density-fitting auxiliary basis. |
None
|
spin
|
int | None
|
2S; defaults to |
None
|
n_radial, lebedev
|
Becke-grid quality (match the energy calculation). |
required |
Minimize the RKS energy directly over orthonormalized coefficients (Adam).
Parameters:
| Name | Type | Description | Default |
|---|---|---|---|
ks
|
RKS
|
the :class: |
required |
learning_rate
|
float
|
Adam step size. |
0.3
|
max_steps
|
int
|
optimizer step budget. |
2000
|
g_tol
|
float
|
stop when the gradient norm falls below this. |
1e-06
|
Z0
|
Float[Array, 'nao nocc'] | None
|
optional initial (nao, nocc) coefficient guess; default is the core-Hamiltonian guess. |
None
|
verbose
|
bool
|
print per-step energy and gradient norm. |
False
|
Minimize the UKS energy directly over per-spin coefficients (Adam).
Args mirror :func:dftax.ks.minimize.rks_minimize; Z0 is an optional
(Zα, Zβ) initial guess (default: per-spin core-Hamiltonian guess).
Batched (vmap over geometries)
Closed-shell RKS over a batch of geometries coords_batch of shape
(B, n_atom, 3) (Bohr). mol supplies the (fixed) atoms + basis; only the
coordinates vary. Exact 4-center ERI path (use small/moderate bases). With
forces=True also returns analytic nuclear forces (B, n_atom, 3).
Response properties
Permanent electric dipole moment μ (atomic units, or Debye if debye).
μ = Σ_A Z_A (R_A − origin) − Tr(P · r) from the converged density. For a
neutral molecule the result is independent of origin.
Static dipole polarizability tensor α_ij = ∂μ_i/∂E_j (atomic units).
method="fd" (default): central finite difference of the dipole under a uniform
field of strength field (robust; one SCF per ±component). method="analytic":
exact coupled-perturbed KS response via implicit differentiation of the SCF fixed
point (:func:~dftax.ks.implicit.implicit_density), a single jax.jacobian
through the converged density, no field stepping. Returned symmetrized.
Nuclear Hessian ∂²E/∂R_A∂R_B (Ha/Bohr², shape (3N, 3N)) by central
finite difference of the analytic Pulay-free forces.
Harmonic vibrational analysis. Returns a dict with frequencies (cm⁻¹,
negative = imaginary), mass-weighted modes (columns), and Cartesian
displacement modes cart_modes (columns, (3N, 3N)).
Harmonic Raman activities (Å⁴/amu, up to the usual constant) from the
polarizability derivatives dα/dQ. Expensive: a polarizability (field FD)
at every ±Cartesian displacement (O(N) Hessians' worth of work).
Alchemical gradient ∂E/∂Z_A at fixed electron count (Ha per unit charge).
Hellmann-Feynman: the converged density is held fixed (Löwdin projector) and the energy is differentiated w.r.t. the nuclear charges, which enter only the nuclear-attraction and nuclear-repulsion terms.
Implicit differentiation
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.
Integrals
Compute the overlap matrix S_μν in the AO basis.
Pure JAX, fully differentiable w.r.t. basis.centers.
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). |
Compute the kinetic energy matrix T_μν in the AO basis.
Pure JAX, fully differentiable w.r.t. basis.centers.
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). |
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.
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 |
Returns:
| Type | Description |
|---|---|
Float[Array, 'nao nao']
|
V matrix, shape (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).
System & grid
A molecular system: atoms, geometry, charge/spin, and a basis name.
nelectron
property
Total electron count (sum of Z minus the total charge).
from_xyz(atom, basis, *, unit='angstrom', charge=0, spin=0)
classmethod
Build from a PySCF-style atom string (Angstrom by default).
atom_coords()
Nuclear coordinates in Bohr, shape (n_atoms, 3).
atom_charges()
Nuclear charges (atomic numbers), shape (n_atoms,).
Build a molecular quadrature grid (points in Bohr, weights).
Each atom contributes n_radial Becke radial shells times a Lebedev
angular grid of lebedev points; the per-atom grids are merged with
Becke's fuzzy-Voronoi partition.
Returns (coords, weights) as JAX arrays of shape (n_grid, 3) and
(n_grid,), differentiable w.r.t. coords_bohr.