xnn.common.models.dispersion_fast.eeq#

The large-regime EEQ operator, evaluated without trigonometry in the solver loop.

EEQSystem applies the Ewald reciprocal sum through structure factors cos(G . r_i), sin(G . r_i) of shape (N, N_G). Above SF_BUDGET they are recomputed at every matrix-vector product, which is 85-90% of a D4 step at 12-24k atoms (2e9 sines and cosines per product at 24k atoms). FastEEQSystem computes the same operator another way:

  • Factorized phases. With G = m_1 b_1 + m_2 b_2 + m_3 b_3 and theta_k = b_k . r, exp(i G . r) = exp(i (m_1 theta_1 + m_2 theta_2)) exp(i m_3 theta_3). The reciprocal vectors are grouped into columns of fixed (m_1, m_2), and the two factors are tabulated once per system: B (C, N) over the columns and P (N, 2 M_3 + 1) over m_3. A_rec v is then two complex matrix products (S = B (v P), the structure factor of v on every column, and B^H (g S) back onto the atoms) and no trigonometry. G and -G contribute equally, so only half of the set is kept, with twice the weight. The tables are formed in float64 from float64 phases, so in float32 they are more accurate than the reference’s float32 phases.

  • Real space. The screened kernel of the neighbor list becomes a sparse CSR matrix (images of the same pair merged), applied with one sparse matrix product.

  • Solves. The charge solve and the constraint solve A y = 1 run in lockstep as two right-hand sides of one conjugate-gradient loop. The residual correction of eeq_charges_large(), which is zero at convergence, stops at the charge solve’s tolerance instead of converging its own (tiny) right-hand side again.

  • Implicit differentiation. The differentiable application of A(r) uses the same factorization on live positions and cell, in column blocks that are recomputed in the backward pass.

The operator is the same matrix, so the charges agree with the reference solve to the solver tolerance.

Module Attributes

FAST_LU_MAX_ATOMS

atoms up to which the fast system factorizes the assembled matrix; conjugate gradients above (the reference switches at 12000 atoms)

TABLE_BUDGET

complex entries of the (C, N) column table kept between products; above it the table is rebuilt in blocks at every product

LIVE_BUDGET

complex entries per column block of the differentiable reciprocal sum

CG_TOL_FLOAT32

the reference's 100 machine epsilons (1.2e-5) leave about 3e-5 e in the charges of a 5k-atom box, where its refined LU solve reaches 2e-6 e; 1e-6 matches that at a few more iterations (EEQReuse uses the same value)

LOW_K_VECTORS

reciprocal vectors (smallest |G|, largest weight) in the low-rank part of the conjugate-gradient preconditioner; 0 for plain Jacobi.

Classes

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

EEQSystem with the fast operator and solves.
Show Source