xnn.common.models.dispersion_fast.eeq.FastEEQSystem#

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

Bases: EEQSystem

EEQSystem with the fast operator and solves.

Takes the arguments of the reference class plus the lattice cell (3, 3) in bohr (periodic case), from which the phase tables are built. Right-hand sides may be (N,) or (N, K).

Preconditioner: the long-range Coulomb modes (the smallest reciprocal vectors, whose weights g ~ 1/G^2 are the largest) are what a Jacobi preconditioner leaves untouched and what slows the conjugate gradients. The periodic system therefore uses M = D + U W U^T: U = [cos, sin] of the LOW_K_VECTORS smallest vectors, W their weights, and D the full diagonal of A without that part; M^-1 is applied through the Woodbury identity (one 2L x 2L Cholesky factor per structure, two N x 2L products per iteration). A preconditioner only changes how fast the iteration converges, not what it converges to.

Parameters:

  • diag (Tensor)

  • rad (Tensor)

  • pos (Tensor)

  • edge_index (Optional[Tensor])

  • edge_vec (Optional[Tensor])

  • alpha (float)

  • gvec (Optional[Tensor])

  • gfac (Optional[Tensor])

  • grid (Optional[Tensor])

  • solver (str)

  • signature (Optional[tuple])

  • cell (Optional[Tensor])

keep_structure_factors = False#

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

precondition(r)[source]#

M^-1 r for r of shape (N,) or (N, K): the matrix-free preconditioner EEQReuse uses for large systems in place of a dense inverse.

Parameters:

r (Tensor)

Return type:

Tensor

matvec(v)[source]#

A v for v of shape (N,) or (N, K) (no autograd graph).

Parameters:

v (Tensor)

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=None, maxiter=1000, atol=None)[source]#

Preconditioned CG for A Y = B, the columns of B in lockstep.

Each column stops at |r| <= tol |b| (or |r| <= atol); a finished column takes no further steps while the others converge. tol defaults to CG_TOL in float64 and CG_TOL_FLOAT32 in float32.

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