xnn.common.models.d4.DFTD4#

class xnn.common.models.d4.DFTD4(s6=1.0, s8=1.20065498, a1=0.40085597, a2=5.02928789, s9=1.0, alp=16.0, ga=3.0, gc=2.0, wf=6.0, cutoff_pair=31.750632654269545, cutoff_triple=21.167088436179696, cutoff_cn=15.875316327134772, cutoff_eeq_cn=13.22943027261231, switch_width_pair=0.0, switch_width_triple=0.0, trainable=False, regime='auto', cutoff_eeq=None, eeq_solver='auto', checkpoint_triplets=True, recompute_pairs=True, triplet_chunk=None, triplet_cache=None, tail_correction=False)[source]#

Bases: Module, FastPathModule

The D4 dispersion model as a TorchScript-compatible energy evaluator.

Holds the damping parameters, the model constants and the reference data, and evaluates the charge- and geometry-dependent dispersion energy of one or several structures from positions, atomic numbers and a neighbor list. Positions are in Angstrom, energies in eV; coordination numbers, charges (e), polarizabilities (bohr^3) and C6 coefficients (hartree bohr^6) are reported in the customary atomic units of the D4 literature.

Parameters:

  • s6 (float, optional) – BJ / ATM damping parameters, by default the PBE0-D4 values of the paper (s6 = 1, s8 = 1.20065498, a1 = 0.40085597, a2 = 5.02928789, s9 = 1, alp = 16). s9 = 0 switches the three-body term off.

  • s8 (float, optional) – BJ / ATM damping parameters, by default the PBE0-D4 values of the paper (s6 = 1, s8 = 1.20065498, a1 = 0.40085597, a2 = 5.02928789, s9 = 1, alp = 16). s9 = 0 switches the three-body term off.

  • a1 (float, optional) – BJ / ATM damping parameters, by default the PBE0-D4 values of the paper (s6 = 1, s8 = 1.20065498, a1 = 0.40085597, a2 = 5.02928789, s9 = 1, alp = 16). s9 = 0 switches the three-body term off.

  • a2 (float, optional) – BJ / ATM damping parameters, by default the PBE0-D4 values of the paper (s6 = 1, s8 = 1.20065498, a1 = 0.40085597, a2 = 5.02928789, s9 = 1, alp = 16). s9 = 0 switches the three-body term off.

  • s9 (float, optional) – BJ / ATM damping parameters, by default the PBE0-D4 values of the paper (s6 = 1, s8 = 1.20065498, a1 = 0.40085597, a2 = 5.02928789, s9 = 1, alp = 16). s9 = 0 switches the three-body term off.

  • alp (float, optional) – BJ / ATM damping parameters, by default the PBE0-D4 values of the paper (s6 = 1, s8 = 1.20065498, a1 = 0.40085597, a2 = 5.02928789, s9 = 1, alp = 16). s9 = 0 switches the three-body term off.

  • ga (float, optional) – Model constants: charge-scaling height beta1 (3.0), charge-scaling steepness multiplier of the chemical hardness (2.0), and the Gaussian weighting exponent beta2 (6.0).

  • gc (float, optional) – Model constants: charge-scaling height beta1 (3.0), charge-scaling steepness multiplier of the chemical hardness (2.0), and the Gaussian weighting exponent beta2 (6.0).

  • wf (float, optional) – Model constants: charge-scaling height beta1 (3.0), charge-scaling steepness multiplier of the chemical hardness (2.0), and the Gaussian weighting exponent beta2 (6.0).

  • cutoff_pair (float, optional) – Real-space cutoffs in Angstrom of the two-body sum, the three-body sum, the D4 coordination number and the EEQ coordination number; by default the upstream 60, 40, 30 and 25 bohr.

  • cutoff_triple (float, optional) – Real-space cutoffs in Angstrom of the two-body sum, the three-body sum, the D4 coordination number and the EEQ coordination number; by default the upstream 60, 40, 30 and 25 bohr.

  • cutoff_cn (float, optional) – Real-space cutoffs in Angstrom of the two-body sum, the three-body sum, the D4 coordination number and the EEQ coordination number; by default the upstream 60, 40, 30 and 25 bohr.

  • cutoff_eeq_cn (float, optional) – Real-space cutoffs in Angstrom of the two-body sum, the three-body sum, the D4 coordination number and the EEQ coordination number; by default the upstream 60, 40, 30 and 25 bohr.

  • switch_width_pair (float, optional) – Widths (Angstrom) of quintic switching windows ending at the pair / triple cutoffs, by default 0 (sharp cutoffs, as upstream). A window of a few Angstrom keeps energy and forces continuous under a finite cutoff, which matters when D4 supplements an MLIP in MD.

  • switch_width_triple (float, optional) – Widths (Angstrom) of quintic switching windows ending at the pair / triple cutoffs, by default 0 (sharp cutoffs, as upstream). A window of a few Angstrom keeps energy and forces continuous under a finite cutoff, which matters when D4 supplements an MLIP in MD.

  • trainable (bool, optional) – Make s6, s8, a1, a2, s9 learnable nn.Parameter``s, by default ``False (fixed buffers).

  • regime (str, optional) – How the EEQ charges are computed: "dense" builds the (N, N) interaction matrix as the reference code does (bit-exact dftd4 parity; memory grows as N^2), "large" uses the matrix-free operator, iterative / factorized solve and implicit differentiation of eeq (memory O(E + N N_G), agreement with dense to about 1e-10 hartree, eager only), and "auto" (default) picks large above AUTO_LARGE_PERIODIC / AUTO_LARGE_MOLECULAR atoms. TorchScript always runs dense.

  • cutoff_eeq (float or None, optional) – Real-space range (Angstrom) of the large-regime Ewald split, which sets the splitting parameter (a longer range means fewer reciprocal vectors). None (default) takes DEFAULT_CUTOFF_EEQ (16 A) or the largest of the other cutoffs, whichever is larger, unless regime="dense" (which has no Ewald split and keeps the other cutoffs’ radius); the neighbor list is then at most 16 A, which costs about 2.4 times the edges of a 12 A list and buys a four times cheaper EEQ operator. The large regime needs at least 20 bohr (10.6 A), below which the erf(gamma r) part of the kernel is not screened, and raises otherwise.

  • eeq_solver (str, optional) – Large-regime solver: "auto" (LU of the assembled matrix up to LU_MAX_ATOMS atoms, conjugate gradients above), "lu" or "cg".

  • checkpoint_triplets (bool, optional) – Evaluate the three-body term in recompute blocks of centers (three_body_energy_chunked(), recompute), by default True: memory is bounded by one block at every derivative order (force training included) and no (N, N) C6 matrix is needed, at the cost of re-evaluating each block once per order of differentiation. Eager only; TorchScript uses the plain chunk loop.

  • recompute_pairs (bool, optional) – Evaluate the two-body term in recompute blocks of edges, by default True: the per-edge (E, 23) polarizability products of the Casimir-Polder C6 (about 1.5 kB per edge when retained for the backward pass, 14 million edges at 20 000 atoms with a 12 A cutoff) are then transient. Eager only.

  • triplet_chunk (int or None, optional) – Triplets per three-body recompute block; None (default) sizes it from the free device memory (2^20 to 2^24). Larger blocks amortize the per-block set-up, which is paid in both passes.

  • triplet_cache (float or None, optional) – Keep each three-body block’s triples from the forward to the backward pass so they are enumerated once per step: a budget in GB, None (default) for a quarter of the free device memory, 0 to disable. Exact (the backward pass uses the same triples); 14 bytes per triple, e.g. 0.5 GB for 5000 water atoms at an 8 A triple cutoff. Eager only.

  • tail_correction (bool, optional) – Add, for periodic structures, the homogeneous long-range correction of the two-body term: the energy the pair cutoff and its switching window remove, E_tail = (2 pi / V) sum_ij int (1 - sw(r)) u_ij(r) r^2 dr over all ordered atom pairs with a uniform distribution beyond the window (the analog of LAMMPS pair_modify tail yes). It uses the structure’s own (charge- and CN-dependent) C6 and the BJ damping, and is differentiable, so its forces (from the C6) and its stress (-E_tail / V on the diagonal) are consistent with the energy. By default False. Molecular structures are unaffected.

Notes

For molecular dynamics and geometry optimization, enable_eeq_reuse() carries the large-regime EEQ solve from one step to the next (see EEQReuse).

Variables:

  • cutoff (float) – The largest of the cutoffs (Angstrom) – the neighbor-list radius this module needs.

  • regime (str) – The requested regime ("auto", "dense" or "large").

Parameters:
coordination_numbers(z, edge_index, r, n_atoms)[source]#

D4 and EEQ coordination numbers of every atom (paper eqs 6, 14).

Parameters:

  • z (Tensor) – Atomic numbers, shape (N,).

  • edge_index (Tensor) – Edge list [src, dst] of shape (2, E); the count of src is added to the center dst.

  • r (Tensor) – Edge lengths in bohr, shape (E,).

  • n_atoms (int) – Number of atoms N.

Returns:

  • cn_d4 (Tensor) – The electronegativity-weighted CN entering the Gaussian weighting, shape (N,) (edges within cutoff_cn).

  • cn_eeq (Tensor) – The plain CN of the charge model, softly capped at 8, shape (N,) (edges within cutoff_eeq_cn).

Return type:

Tuple[Tensor, Tensor]

eeq_charges(z, pos, cn_eeq, total_charge, cell, periodic)[source]#

EEQ partial charges of one structure (paper eqs 11-16).

Parameters:

  • z (Tensor) – Atomic numbers, shape (N,).

  • pos (Tensor) – Positions in bohr, shape (N, 3).

  • cn_eeq (Tensor) – Capped error-function coordination numbers, shape (N,).

  • total_charge (Tensor) – Scalar total charge of the structure.

  • cell (Tensor) – Lattice vectors as rows in bohr, shape (3, 3) (ignored when periodic is False).

  • periodic (bool) – Whether to Ewald-sum the Coulomb matrix.

Returns:

Partial charges, shape (N,), summing to total_charge.

Return type:

Tensor

fast_supported(device, dtype)[source]#

The fast paths (dispersion_fast) run on CUDA.

Parameters:
Return type:

bool

select_regime(n_atoms, periodic)[source]#

Resolve "auto" to "dense" or "large" for one structure.

Parameters:
Return type:

str

enable_eeq_reuse(enabled=True, **options)[source]#

Carry the large-regime EEQ solve over between consecutive structures.

For molecular dynamics and geometry optimization, where one structure follows the previous one closely; see EEQReuse, which receives options (tol, refresh, maxiter). Results agree with the fresh solve to the solver tolerance; the dense regime is not affected. enable_eeq_reuse(False) switches it off again. Eager only.

Parameters:

enabled (bool)

Return type:

None

reference_weights(z, cn, q)[source]#

Charge-scaled Gaussian weights of the reference systems (eqs 2-4, 8).

Parameters:

  • z (Tensor) – Atomic numbers, shape (N,).

  • cn (Tensor) – D4 coordination numbers, shape (N,).

  • q (Tensor) – Partial charges, shape (N,) (zeros for the ATM term).

Returns:

W_A,ref zeta(z_A, z_A,ref) for the 7 reference slots of each atom, shape (N, 7) (zero in unused slots).

Return type:

Tensor

dynamic_polarizabilities(z, weights)[source]#

Atom-in-molecule alpha_A(i omega) on the 23-point grid (eq 7).

Returns shape (N, 23) in bohr^3; column 0 (omega = 1e-6) is the static polarizability.

Parameters:
Return type:

Tensor

pair_c6(alpha_i, alpha_j)[source]#

C6 of pairs from their dynamic polarizabilities (eq 9), (E,).

Parameters:
Return type:

Tensor

two_body_energy(z, edge_index, r, alpha_iw, n_atoms)[source]#

Per-atom BJ-damped two-body energy in hartree (eqs 18-21).

Every directed edge contributes half of its pair energy to its center, so the sum over atoms is the sum over unordered pairs.

Parameters:
Return type:

Tensor

tail_integrals()[source]#

Radial integrals of the two-body term lost to the pair cutoff.

J_AB = -int_{r_s}^inf (1 - sw(r)) [s6 / (r^6 + R0^6) + s8 (C8/C6) / (r^8 + R0^8)] r^2 dr per element pair, in bohr^-3 (multiply by a C6 in hartree bohr^6), with r_s the start of the switching window (the cutoff itself for a sharp one). The window is integrated with Gauss-Legendre in r, the region beyond the cutoff with Gauss-Legendre in x = 1 / r.

Returns:

Shape (Z_max + 1, Z_max + 1).

Return type:

Tensor

tail_energy(z, alpha_iw, cell)[source]#

Per-atom homogeneous tail correction of one periodic structure (hartree).

E_tail = (2 pi / V) sum_AB J_AB S_AB with S_AB the sum of C6 over the ordered pairs of atoms of elements A and B, which the Casimir-Polder form reduces to per-element sums of the dynamic polarizabilities. Each atom carries its pairs with every other atom.

Parameters:

  • z (Tensor) – Atomic numbers, shape (N,).

  • alpha_iw (Tensor) – Dynamic polarizabilities, shape (N, 23).

  • cell (Tensor) – Lattice vectors as rows in bohr, shape (3, 3).

Returns:

Shape (N,).

Return type:

Tensor

pair_radius_table()[source]#

BJ critical radii a1 sqrt(3 Q_A Q_B) + a2 for all element pairs, bohr.

Return type:

Tensor

evaluate(atomic_numbers, pos, edge_index, edge_vec, batch, num_graphs, cell, pbc, total_charge)[source]#

Dispersion energy and D4 properties of a batch of structures.

The TorchScript-compatible core shared by forward() and the deploy wrappers. Inputs in Angstrom; the neighbor list must reach cutoff (edges beyond a term’s own cutoff are ignored by it).

Parameters:

  • atomic_numbers (Tensor) – Atomic numbers, shape (N,).

  • pos (Tensor) – Cartesian positions in Angstrom, shape (N, 3).

  • edge_index (Tensor) – Edge list [src, dst], shape (2, E), both directions of every pair present.

  • edge_vec (Tensor) – Edge vectors pos[dst] - pos[src] (+ shift) in Angstrom, shape (E, 3).

  • batch (Tensor) – Structure index of every atom, shape (N,).

  • num_graphs (int) – Number of structures B.

  • cell (Tensor) – Lattice vectors as rows in Angstrom, shape (B, 3, 3); an all-zero cell marks a molecular structure.

  • pbc (Tensor) – Periodicity flags, shape (B, 3); a structure is periodic when any flag is set and its cell is nonzero.

  • total_charge (Tensor) – Total charge per structure, shape (B,).

Returns:

"node_energy" (N,) in eV, plus the atomic units quantities "coordination_numbers" (N,), "charges" (N,), "polarizabilities" (N,) (static, bohr^3) and "dynamic_polarizabilities" (N, 23) (for c6_matrix()), and "energy_2body" / "energy_3body" (B,) in eV.

Return type:

dict of str to Tensor

forward(data)[source]#

Evaluate a (batched) AtomicGraph.

Adds "energy" (B,) to the keys of evaluate().

Return type:

Dict[str, Tensor]