xnn.common.models.eeq.EEQReuse#

class xnn.common.models.eeq.EEQReuse(tol=None, refresh=15, maxiter=200)[source]#

Bases: object

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

Between the steps of molecular dynamics (or of a geometry optimization) the atoms move by hundredths of an Angstrom and the EEQ matrix hardly changes, so the previous step’s solve is worth reusing instead of assembling and factorizing the matrix every step:

  • the preconditioner is the explicit inverse of an earlier step’s matrix, formed once and re-formed only when a solve needs more than refresh iterations; applying it is one matrix-vector product. Above DENSE_REUSE_MAX_ATOMS atoms, a system that brings its own matrix-free preconditioner (the fast path’s) is preconditioned with that instead, so no (N, N) matrix is formed;

  • each of the step’s solves is conjugate gradients from a good initial guess: the charges from a quadratic extrapolation of the last three steps, the adjoint (forces) from the previous step’s adjoint, A^-1 1 (the charge constraint) from the previous step’s; the residual correction of eeq_charges_large() uses an absolute tolerance scaled to the charge right-hand side, so it costs nothing when the first solve was already converged.

The results do not depend on the reuse beyond the tolerance tol (relative residual; by default 1e-9 in float64 and 1e-6 in float32): a poor preconditioner or guess only costs iterations. If a solve does not converge within maxiter the preconditioner is re-formed from the current matrix and the solve retried once, and failing that the step falls back to the LU path (a cold iterative solve for a matrix-free system).

Measured on 5001 water atoms at 0.997 g/cm^3 (A100, cutoff_eeq 16 A, a real 0.5 fs trajectory): the three solves of a step take 20 ms in float32 and 59 ms in float64, against 0.13-0.31 s for assembly + LU, with one inverse formed for the whole run.

One instance serves one sequence of structures of the same system (atom count, elements, dtype, device); it resets itself when any of these change, and DFTD4 uses it only for a structure evaluated on its own, never inside a batch. With the dense preconditioner it holds an (N, N) matrix (0.1 GB at 5000 atoms in float32).

Parameters:

  • tol (float or None) – Relative residual of every solve; None picks it from the dtype.

  • refresh (int) – Re-form the preconditioner for the next step when a step’s solves need more than this many iterations in total.

  • maxiter (int) – Iterations after which a solve is considered failed.

reset()[source]#

Forget the previous structures (a new run, or a different system).

Return type:

None

solve(system, b_top, b_bot, role='charges')[source]#

The constrained solve of system for one role of the step ("charges", "residual", "adjoint"; see EEQSystem.solve_augmented()).

Parameters:
Return type:

tuple[Tensor, Tensor]