xnn.common.models.eeq.EEQSystem#

class xnn.common.models.eeq.EEQSystem(diag, rad, pos, edge_index=None, edge_vec=None, alpha=0.0, gvec=None, gfac=None, grid=None, solver='auto', reuse=None, signature=None)[source]#

Bases: object

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

Built from detached geometry: it is the fixed Jacobian of the implicit differentiation. apply_differentiable() re-evaluates the same operator on live tensors for the residual.

Parameters:

  • diag (Tensor) – Diagonal J_A + sqrt(2/pi) / a_A (periodic: - 2 alpha / sqrt(pi)), (N,).

  • rad (Tensor) – EEQ charge widths a_A in bohr, (N,).

  • pos (Tensor) – Positions in bohr, (N, 3); molecular kernel and structure factors.

  • edge_index (Tensor or None) – Directed neighbor list within the real-space cutoff (periodic case), (2, E) and (E, 3) in bohr; None for a molecule.

  • edge_vec (Tensor or None) – Directed neighbor list within the real-space cutoff (periodic case), (2, E) and (E, 3) in bohr; None for a molecule.

  • alpha (float) – Ewald splitting parameter (periodic case), 0 for a molecule.

  • gvec (Tensor or None) – Reciprocal vectors (N_G, 3), weights (N_G,) and their integer triples (periodic), from reciprocal_vectors().

  • gfac (Tensor or None) – Reciprocal vectors (N_G, 3), weights (N_G,) and their integer triples (periodic), from reciprocal_vectors().

  • grid (Tensor or None) – Reciprocal vectors (N_G, 3), weights (N_G,) and their integer triples (periodic), from reciprocal_vectors().

  • solver (str) – "auto" (LU up to LU_MAX_ATOMS, else CG), "lu" or "cg".

  • reuse (EEQReuse or None) – Carry the solve over from the previous structure of an MD run or optimization (EEQReuse); solver is then only the fallback.

  • signature (tuple or None) – Hashable identity of the system (atom count and elements) that EEQReuse compares before reusing its state.

keep_structure_factors = True#

keep the (N, N_G) cos / sin tables when they fit in SF_BUDGET (a subclass with its own reciprocal sum turns this off)

matvec(v)[source]#

A v for the detached geometry, (N,) (no autograd graph).

Parameters:

v (Tensor)

Return type:

Tensor

assemble()[source]#

The dense (N, N) matrix of the operator (for the LU path and tests).

Return type:

Tensor

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

Solve [[A, 1], [1^T, 0]] [y; mu] = [b_top; b_bot]; returns (y, mu).

role names the solve of the step: "charges", "residual" (the correction of eeq_charges_large()) or "adjoint" (the backward pass); EEQReuse keeps a separate history per role.

Parameters:
Return type:

tuple[Tensor, Tensor]

conjugate_gradient(b, tol=1e-12, maxiter=1000)[source]#

Jacobi-preconditioned CG for A y = b (A symmetric positive definite).

Parameters:
Return type:

Tensor

apply_differentiable(pos, edge_vec, rad, diag, q, cell=None)[source]#

A(r) q on live tensors in recomputed blocks (memory O(E + N c) at any order).

cell is the live lattice (periodic case): the reciprocal vectors and weights are rebuilt from it so strain derivatives (stress) flow.

Parameters:
Return type:

Tensor