xnn.common.models.eeq#

Large-system EEQ charges: matrix-free operator, solvers, implicit differentiation.

The dense EEQ path of DFTD4 builds the (N, N) interaction matrix with autograd recording every pairwise intermediate (about eight (N, N) tensors for a molecule, and for a periodic cell the Wigner-Seitz image search plus 125 translations per pair, about 7 kB per pair). It reproduces dftd4 bit for bit and is the right tool up to a few thousand atoms. This module is the large regime: the same linear system, evaluated so that memory stays O(E + N N_G + N c) (edges, structure factors, one row block) and gradients cost one matrix-vector product per backward.

Algorithm:

  • Operator. Molecules: A_ij = erf(gamma_ij r_ij) / r_ij applied in row blocks that are never stored. Periodic cells: an Ewald split at a splitting parameter alpha chosen from the real-space cutoff, the screened kernel [erf(gamma r) - erf(alpha r)] / r summed over the neighbor list (all images within the cutoff, self-images included), and the reciprocal sum written through structure factors: with C = cos(r . G), S = sin(r . G) of shape (N, N_G) the reciprocal block is the low-rank product C diag(g) C^T + S diag(g) S^T, so A v costs O(N N_G) without forming (N, N). The diagonal is J + sqrt(2/pi) / a (minus 2 alpha / sqrt(pi) in the periodic case), as in the reference code.

  • Solve. The charge-constrained system [[A, 1], [1^T, 0]] [q; mu] = [x; Q] is solved with an LU factorization of the assembled matrix while (N, N) fits, and otherwise by conjugate gradients on the symmetric positive definite A (Jacobi preconditioner) with the constraint eliminated: q = y_1 - mu y_2, A y_1 = x, A y_2 = 1.

  • Gradients. Implicit differentiation with a constant Jacobian: the solution q_0 is computed without autograd and refined once, q_1 = q_0 + J^{-1} [x - A q_0 - mu_0; Q - 1^T q_0] (also without autograd), then q = q_1 + J^{-1} [x - A(r) q_1 - mu_1; Q - 1^T q_1] with the constant operator J^{-1} as a custom autograd function whose value is taken as zero (it is, to rounding: q_1 is converged) and whose backward is another solve, around one differentiable application of A(r) in recomputed blocks. Its derivative is the exact first-order sensitivity dq/dr = -J^{-1} (dA/dr) q at q_1: expanding around the refined charges matters because an iterative q_0 carries a relative error of the solver tolerance times the condition number, which the refinement removes from the charges and the expansion point removes from the forces (expanded around q_0 the forces were off by ~1e-10 hartree/bohr at a 1e-12 tolerance). Force training (a second derivative with respect to the model parameters, on which A does not depend) is exact.

The large regime agrees with the dense one to the truncation and solver tolerances (about 1e-10 hartree), not to the last bit, and it is eager-only: none of this scripts, so TorchScript exports keep the dense path.

Module Attributes

EWALD_EPS

relative truncation of the real- and reciprocal-space Ewald tails

CG_TOL

conjugate-gradient stopping criterion, |r| <= tol |b| (raised to 100 machine epsilons in single precision)

LU_MAX_ATOMS

atoms up to which the assembled matrix is LU-factorized (two (N, N) float64 tensors: 2.3 GB at 12000 atoms)

ROW_CHUNK

ROW_CHUNK x N)

SF_BUDGET

3.2 GB); beyond it the reciprocal sum is evaluated in chunks of G vectors

DENSE_REUSE_MAX_ATOMS

above this many atoms a system with its own matrix-free preconditioner (the fast path's low-rank one) is preconditioned with it; up to it, the explicit inverse of an earlier step's matrix (at 10000 atoms 0.8 GB in float64, 2.4 GB while it is formed)

Functions

eeq_charges_large(system, pos, edge_vec, ...)

Constrained EEQ charges with implicit differentiation, (N,).

ewald_alpha(cutoff[, eps])

Splitting parameter with erfc(alpha r_c) = eps at the real-space cutoff (bohr).

lu_solve_implicit(lu, pivots, full, rhs, rounds)

full^-1 @ rhs with a fixed (detached) LU factor, refined, and differentiable through the residual instead of through the factor.

lu_solve_refined(lu, pivots, full, rhs, rounds)

full^-1 @ rhs from the LU factor of full, with rounds rounds of iterative refinement.

reciprocal_vectors(cell, alpha[, eps])

Reciprocal vectors G != 0 with exp(-G^2 / 4 alpha^2) >= eps and their weights.

reciprocal_weights(grid, cell, alpha)

G = m B and g(G) for integer triples grid and a (live) cell.

Classes

EEQReuse([tol, refresh, maxiter])

Carry the EEQ solve from one structure to the next of an MD run or optimization.

EEQSystem(diag, rad, pos[, edge_index, ...])

The EEQ linear system of one structure as a constant, matrix-free operator.
Show Source