Skip to content

Commit

Permalink
Restore compatibility between lr_eig and real_eig
Browse files Browse the repository at this point in the history
  • Loading branch information
@sunqm
sunqm committed Nov 21, 2024
1 parent 3ff60cc commit 5766764
Show file tree
Hide file tree
Showing 5 changed files with 63 additions and 169 deletions.
47 changes: 26 additions & 21 deletions pyscf/tdscf/_lr_eig.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -89,9 +89,9 @@ def eigh(aop, x0, precond, tol_residual=1e-5, lindep=1e-12, nroots=1,
space_inc = max(nroots, min(MAX_SPACE_INC, x0_size//2))

max_space = int(max_memory*1e6/8/x0_size / 2 - nroots - space_inc)
if max_space < nroots * 2 < x0_size:
if max_space < nroots * 4 < x0_size:
log.warn('Not enough memory to store trial space in _lr_eig.eigh')
max_space = max(max_space, nroots * 2)
max_space = max(max_space, nroots * 4)
max_space = min(max_space, x0_size)
log.debug(f'Set max_space {max_space}, space_inc {space_inc}')

Expand Down Expand Up @@ -271,6 +271,8 @@ def eig(aop, x0, precond, tol_residual=1e-5, nroots=1, x0sym=None, pick=None,
Allowed memory in MB.
Returns:
conv : list of booleans
Whether each root is converged.
e : list of floats
Eigenvalues.
c : list of 1D arrays
Expand All @@ -296,10 +298,10 @@ def eig(aop, x0, precond, tol_residual=1e-5, nroots=1, x0sym=None, pick=None,
space_inc = max(nroots, min(MAX_SPACE_INC, x0_size//4))

max_space = int(max_memory*1e6/8/(2*x0_size) / 2 - space_inc)
if max_space < nroots * 2 < x0_size:
if max_space < nroots * 4 < x0_size:
log.warn('Not enough memory to store trial space in _lr_eig.eig')
max_space = space_inc * 2
max_space = max(max_space, nroots * 2)
max_space = max(max_space, nroots * 4)
max_space = min(max_space, x0_size)
log.debug(f'Set max_space {max_space}, space_inc {space_inc}')

Expand Down Expand Up @@ -473,7 +475,7 @@ def eig(aop, x0, precond, tol_residual=1e-5, nroots=1, x0sym=None, pick=None,

return conv[:nroots], e[:nroots], x0[:nroots]

def real_eig(matvec, x0, precond, tol_residual=1e-5, nroots=1, x0sym=None, pick=None,
def real_eig(aop, x0, precond, tol_residual=1e-5, nroots=1, x0sym=None, pick=None,
max_cycle=50, max_memory=MAX_MEMORY, lindep=1e-12, verbose=logger.WARN):
'''
Solve linear response eigenvalues for real A and B matrices
Expand All @@ -484,14 +486,11 @@ def real_eig(matvec, x0, precond, tol_residual=1e-5, nroots=1, x0sym=None, pick=
algorithm implemented in https://github.com/John-zzh/improved-Davidson-Algorithm
Args:
matvec : function(x, y) => (array_like_x, array_like_y)
The matrix-vector product operator to perform.
[ A B] [X]
[-B -A] [Y]
The shape of x is [nroots,nocc,nvir].
x0 : (X_array, Y_array)
aop : function(x) => array_like_x
The matrix-vector product operator.
x0 : 1D array or a list of 1D array
Initial guess.
precond : function(x, y, e) => (array_like_x, array_like_y)
precond : function(dx, e) => array_like_dx
Preconditioner to generate new trial vector.
Kwargs:
Expand All @@ -515,7 +514,7 @@ def real_eig(matvec, x0, precond, tol_residual=1e-5, nroots=1, x0sym=None, pick=
Whether each root is converged.
e : list of floats
Eigenvalues.
c : (X_array, Y_array)
c : list of 1D arrays
Eigenvectors.
'''

Expand All @@ -527,20 +526,22 @@ def real_eig(matvec, x0, precond, tol_residual=1e-5, nroots=1, x0sym=None, pick=
else:
log = logger.Logger(sys.stdout, verbose)

V, W = x0
assert V.ndim == 2 and W.ndim == 2
A_size = V.shape[1]
assert x0.ndim == 2
A_size = x0.shape[1] // 2
V = x0[:,:A_size]
W = x0[:,A_size:]
x0 = (V, W)
if MAX_SPACE_INC is None:
space_inc = nroots
else:
# Adding too many trial bases in each iteration may cause larger errors
space_inc = max(nroots, min(MAX_SPACE_INC, A_size//2))

max_space = int(max_memory*1e6/8/(4*A_size) / 2 - space_inc)
if max_space < nroots * 2 < A_size:
if max_space < nroots * 4 < A_size:
log.warn('Not enough memory to store trial space in _lr_eig.eig')
max_space = space_inc * 2
max_space = max(max_space, nroots * 2)
max_space = max(max_space, nroots * 4)
max_space = min(max_space, A_size)
log.debug(f'Set max_space {max_space}, space_inc {space_inc}')

Expand Down Expand Up @@ -571,7 +572,9 @@ def real_eig(matvec, x0, precond, tol_residual=1e-5, nroots=1, x0sym=None, pick=
if x0sym is not None:
xs_ir = xt_ir = x0_ir

U1, U2 = matvec(V, W)
axt = aop(np.hstack([V, W]))
U1 = axt[:,:A_size]
U2 = -axt[:,A_size:]
m0, m1 = m1, m1+len(U1)
V_holder [:,m0:m1] = V.T
W_holder [:,m0:m1] = W.T
Expand Down Expand Up @@ -689,7 +692,9 @@ def real_eig(matvec, x0, precond, tol_residual=1e-5, nroots=1, x0sym=None, pick=
break

r_index = r_norms > tol_residual
X_new, Y_new = precond(R_x[:,r_index], R_y[:,r_index], w[r_index])
XY_new = precond(np.vstack([R_x[:,r_index], R_y[:,r_index]]).T, w[r_index])
X_new = XY_new[:,:A_size].T
Y_new = XY_new[:,A_size:].T
if x0sym is None:
V, W = VW_Gram_Schmidt_fill_holder(
V_holder[:,:m1], W_holder[:,:m1], X_new, Y_new)
Expand Down Expand Up @@ -731,7 +736,7 @@ def real_eig(matvec, x0, precond, tol_residual=1e-5, nroots=1, x0sym=None, pick=
if len(x0[0]) < min(A_size, nroots):
log.warn(f'Not enough eigenvectors (len(x0)={len(x0[0])}, nroots={nroots})')

return conv[:nroots], e[:nroots], x0
return conv[:nroots], e[:nroots], np.hstack(x0)

def _gen_x0(v, xs):
out = _outprod_to_subspace(v[::2], xs)
Expand Down
103 changes: 19 additions & 84 deletions pyscf/tdscf/rhf.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -748,6 +748,8 @@ def get_ab(self, mf=None):

def get_precond(self, hdiag):
def precond(x, e, *args):
if isinstance(e, numpy.ndarray):
e = e[0]
diagd = hdiag - (e-self.level_shift)
diagd[abs(diagd)<1e-8] = 1e-8
return x/diagd
Expand Down Expand Up @@ -916,8 +918,7 @@ def pickeig(w, v, nroots, envs):
CIS = TDA


def gen_tdhf_operation(mf, fock_ao=None, singlet=True, wfnsym=None,
real_eig_solver=True):
def gen_tdhf_operation(mf, fock_ao=None, singlet=True, wfnsym=None):
'''Generate function to compute
[ A B ][X]
Expand Down Expand Up @@ -952,37 +953,6 @@ def gen_tdhf_operation(mf, fock_ao=None, singlet=True, wfnsym=None,
max_memory = max(2000, mf.max_memory*.8-mem_now)
vresp = mf.gen_response(singlet=singlet, hermi=0, max_memory=max_memory)

if real_eig_solver:
def vind(xs, ys):
nz = len(xs)
xs = numpy.asarray(xs).reshape(nz,nocc,nvir)
ys = numpy.asarray(ys).reshape(nz,nocc,nvir)
if wfnsym is not None and mol.symmetry:
xs[:,sym_forbid] = 0
ys[:,sym_forbid] = 0

# *2 for double occupancy
dms = lib.einsum('xov,pv,qo->xpq', xs, orbv, orbo.conj()*2)
dms += lib.einsum('xov,qv,po->xpq', ys, orbv.conj(), orbo*2)
v1ao = vresp(dms) # = <mj||nb> Xjb + <mb||nj> Yjb
# A ~= <aj||ib>, B = <ab||ij>
# AX + BY
# = <aj||ib> Xjb + <ab||ij> Yjb
# = (<mj||nb> Xjb + <mb||nj> Yjb) Cma* Cni
v1_top = lib.einsum('xpq,qo,pv->xov', v1ao, orbo, orbv.conj())
# (B*)X + (A*)Y
# = <ij||ab> Xjb + <ib||aj> Yjb
# = (<mj||nb> Xjb + <mb||nj> Yjb) Cmi* Cna
v1_bot = lib.einsum('xpq,po,qv->xov', v1ao, orbo.conj(), orbv)
v1_top += numpy.einsum('xia,ia->xia', xs, e_ia) # AX
v1_bot += numpy.einsum('xia,ia->xia', ys, e_ia) # (A*)Y

if wfnsym is not None and mol.symmetry:
v1_top[:,sym_forbid] = 0
v1_bot[:,sym_forbid] = 0
return v1_top.reshape(nz,nocc*nvir), v1_bot.reshape(nz,nocc*nvir)
return vind, hdiag.ravel()

def vind(xys):
xys = numpy.asarray(xys).reshape(-1,2,nocc,nvir)
if wfnsym is not None and mol.symmetry:
Expand Down Expand Up @@ -1044,47 +1014,19 @@ class TDHF(TDBase):
'''

@lib.with_doc(gen_tdhf_operation.__doc__)
def gen_vind(self, mf=None, real_eig_solver=True):
def gen_vind(self, mf=None):
if mf is None:
mf = self._scf
return gen_tdhf_operation(mf, None, self.singlet, self.wfnsym,
real_eig_solver)

def get_precond(self, hdiag, real_eig_solver=True):
if real_eig_solver:
def precond(x, y, e):
'''preconditioners for each corresponding residual (state)'''
diagd = hdiag - (e[0]-self.level_shift)
diagd[abs(diagd)<1e-8] = 1e-8
X_new = x / diagd[:,None]
Y_new = y / diagd[:,None]
return X_new, Y_new
else:
def precond(x, e, *args):
diagd = hdiag - (e-self.level_shift)
diagd[abs(diagd)<1e-8] = 1e-8
return x/diagd
return precond
return gen_tdhf_operation(mf, None, self.singlet, self.wfnsym)

def init_guess(self, mf, nstates=None, wfnsym=None, return_symmetry=False,
real_eig_solver=True):
'''
Kwargs:
real_eig_solver : bool
Create initial guess for _lr_eig.real_eig . It contains two
arrays, X and Y, than a single vector.
'''
def init_guess(self, mf, nstates=None, wfnsym=None, return_symmetry=False):
if return_symmetry:
x0, x0sym = TDA.init_guess(self, mf, nstates, wfnsym, return_symmetry)
y0 = numpy.zeros_like(x0)
if real_eig_solver:
return (x0, y0), x0sym
return numpy.hstack([x0, y0]), x0sym
else:
x0 = TDA.init_guess(self, mf, nstates, wfnsym, return_symmetry)
y0 = numpy.zeros_like(x0)
if real_eig_solver:
return x0, y0
return numpy.hstack([x0, y0])

def kernel(self, x0=None, nstates=None):
Expand All @@ -1108,11 +1050,13 @@ def kernel(self, x0=None, nstates=None):

real_eig_solver = real_system

vind, hdiag = self.gen_vind(self._scf, real_eig_solver)
precond = self.get_precond(hdiag, real_eig_solver)
vind, hdiag = self.gen_vind(self._scf)
precond = self.get_precond(hdiag)
if real_eig_solver:
eig = real_eig
pickeig = None
else:
eig = lr_eig
# We only need positive eigenvalues
def pickeig(w, v, nroots, envs):
realidx = numpy.where((abs(w.imag) < REAL_EIG_THRESHOLD) &
Expand All @@ -1125,38 +1069,29 @@ def pickeig(w, v, nroots, envs):
x0sym = None
if x0 is None:
x0, x0sym = self.init_guess(
self._scf, self.nstates, return_symmetry=True,
real_eig_solver=real_eig_solver)
self._scf, self.nstates, return_symmetry=True)
elif mol.symmetry:
x_sym = y_sym = _get_x_sym_table(self._scf).ravel()
x_sym = numpy.append(x_sym, y_sym)
x0sym = [_guess_wfnsym_id(self, x_sym, x) for x in x0]

self.converged, self.e, x1 = eig(
vind, x0, precond, tol_residual=self.conv_tol, lindep=self.lindep,
nroots=nstates, x0sym=x0sym, pick=pickeig, max_cycle=self.max_cycle,
max_memory=self.max_memory, verbose=log)

nocc = numpy.count_nonzero(self._scf.mo_occ)
nmo = self._scf.mo_occ.size
nvir = nmo - nocc

if real_eig_solver:
self.converged, self.e, x1 = real_eig(
vind, x0, precond, tol_residual=self.conv_tol, lindep=self.lindep,
nroots=nstates, x0sym=x0sym, pick=pickeig, max_cycle=self.max_cycle,
max_memory=self.max_memory, verbose=log)
x1 = zip(*x1)

else:
self.converged, self.e, x1 = lr_eig(
vind, x0, precond, tol_residual=self.conv_tol, lindep=self.lindep,
nroots=nstates, x0sym=x0sym, pick=pickeig, max_cycle=self.max_cycle,
max_memory=self.max_memory, verbose=log)
x1 = [z.reshape(2,-1) for z in x1]

def norm_xy(x, y):
def norm_xy(z):
x, y = z.reshape(2, -1)
norm = lib.norm(x)**2 - lib.norm(y)**2
if norm < 0:
log.warn('TDDFT amplitudes |X| smaller than |Y|')
norm = abs(.5/norm) ** .5 # normalize to 0.5 for alpha spin
return x.reshape(nocc,nvir)*norm, y.reshape(nocc,nvir)*norm
self.xy = [norm_xy(x, y) for x, y in x1]
self.xy = [norm_xy(z) for z in x1]

if self.chkfile:
lib.chkfile.save(self.chkfile, 'tddft/e', self.e)
Expand Down
1 change: 1 addition & 0 deletions pyscf/tdscf/rks.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -47,6 +47,7 @@ class CasidaTDDFT(TDDFT, TDA):
'''

init_guess = TDA.init_guess
get_precond = TDA.get_precond

def gen_vind(self, mf=None):
if mf is None:
Expand Down
Loading

0 comments on commit 5766764

Please sign in to comment.