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)[source]#

Bases: Module

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

Variables:

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

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

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

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]