From 5766764b75d94b0ca6b9bbf946d72bef118536bc Mon Sep 17 00:00:00 2001 From: Qiming Sun Date: Thu, 21 Nov 2024 06:22:23 -0800 Subject: [PATCH] Restore compatibility between lr_eig and real_eig --- pyscf/tdscf/_lr_eig.py | 47 ++++++++++--------- pyscf/tdscf/rhf.py | 103 ++++++++--------------------------------- pyscf/tdscf/rks.py | 1 + pyscf/tdscf/uhf.py | 80 +++++++------------------------- pyscf/tdscf/uks.py | 1 + 5 files changed, 63 insertions(+), 169 deletions(-) diff --git a/pyscf/tdscf/_lr_eig.py b/pyscf/tdscf/_lr_eig.py index 445455bb6a..bb307bb6f6 100644 --- a/pyscf/tdscf/_lr_eig.py +++ b/pyscf/tdscf/_lr_eig.py @@ -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}') @@ -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 @@ -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}') @@ -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 @@ -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: @@ -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. ''' @@ -527,9 +526,11 @@ 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: @@ -537,10 +538,10 @@ def real_eig(matvec, x0, precond, tol_residual=1e-5, nroots=1, x0sym=None, pick= 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}') @@ -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 @@ -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) @@ -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) diff --git a/pyscf/tdscf/rhf.py b/pyscf/tdscf/rhf.py index 1f69234a24..9e85ae1f6a 100644 --- a/pyscf/tdscf/rhf.py +++ b/pyscf/tdscf/rhf.py @@ -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 @@ -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] @@ -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) # = Xjb + Yjb - # A ~= , B = - # AX + BY - # = Xjb + Yjb - # = ( Xjb + Yjb) Cma* Cni - v1_top = lib.einsum('xpq,qo,pv->xov', v1ao, orbo, orbv.conj()) - # (B*)X + (A*)Y - # = Xjb + Yjb - # = ( Xjb + 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: @@ -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): @@ -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) & @@ -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) diff --git a/pyscf/tdscf/rks.py b/pyscf/tdscf/rks.py index 5b22e06449..d840a4c50b 100644 --- a/pyscf/tdscf/rks.py +++ b/pyscf/tdscf/rks.py @@ -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: diff --git a/pyscf/tdscf/uhf.py b/pyscf/tdscf/uhf.py index 7eaee702c4..70a994162f 100644 --- a/pyscf/tdscf/uhf.py +++ b/pyscf/tdscf/uhf.py @@ -731,8 +731,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] @@ -773,39 +772,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(hermi=0, max_memory=max_memory) - if real_eig_solver: - def vind(xs, ys): - nz = len(xs) - if wfnsym is not None and mol.symmetry: - xs[:,sym_forbid] = 0 - ys[:,sym_forbid] = 0 - - xa = xs[:,:nocca*nvira].reshape(nz,nocca,nvira) - xb = xs[:,nocca*nvira:].reshape(nz,noccb,nvirb) - ya = ys[:,:nocca*nvira].reshape(nz,nocca,nvira) - yb = ys[:,nocca*nvira:].reshape(nz,noccb,nvirb) - dmsa = lib.einsum('xov,pv,qo->xpq', xa, orbva, orboa.conj()) - dmsb = lib.einsum('xov,pv,qo->xpq', xb, orbvb, orbob.conj()) - dmsa += lib.einsum('xov,qv,po->xpq', ya, orbva.conj(), orboa) - dmsb += lib.einsum('xov,qv,po->xpq', yb, orbvb.conj(), orbob) - v1ao = vresp(numpy.asarray((dmsa,dmsb))) - v1a_top = lib.einsum('xpq,qo,pv->xov', v1ao[0], orboa, orbva.conj()) - v1b_top = lib.einsum('xpq,qo,pv->xov', v1ao[1], orbob, orbvb.conj()) - v1a_bot = lib.einsum('xpq,po,qv->xov', v1ao[0], orboa.conj(), orbva) - v1b_bot = lib.einsum('xpq,po,qv->xov', v1ao[1], orbob.conj(), orbvb) - - v1_top = xs * e_ia # AX - v1_bot = ys * e_ia # AY - v1_top[:,:nocca*nvira] += v1a_top.reshape(nz,-1) - v1_bot[:,:nocca*nvira] += v1a_bot.reshape(nz,-1) - v1_top[:,nocca*nvira:] += v1b_top.reshape(nz,-1) - v1_bot[:,nocca*nvira:] += v1b_bot.reshape(nz,-1) - if wfnsym is not None and mol.symmetry: - v1_top[:,sym_forbid] = 0 - v1_bot[:,sym_forbid] = 0 - return v1_top, v1_bot - return vind, hdiag - def vind(xys): nz = len(xys) xys = numpy.asarray(xys).reshape(nz,2,-1) @@ -850,29 +816,21 @@ class TDHF(TDBase): singlet = None @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) + 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): + 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]) - get_precond = rhf.TDHF.get_precond - def kernel(self, x0=None, nstates=None): '''TDHF diagonalization with non-Hermitian eigenvalue solver ''' @@ -895,11 +853,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) & @@ -909,33 +869,25 @@ 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_a, x_sym_b = _get_x_sym_table(self._scf) x_sym = y_sym = numpy.append(x_sym_a.ravel(), x_sym_b.ravel()) x_sym = numpy.append(x_sym, y_sym) x0sym = [rhf._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) + nmo = self._scf.mo_occ[0].size nocca, noccb = self._scf.nelec nvira = nmo - nocca nvirb = nmo - noccb - 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] - xy = [] - for i, (x, y) in enumerate(x1): + for i, z in enumerate(x1): + 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|') diff --git a/pyscf/tdscf/uks.py b/pyscf/tdscf/uks.py index 25c49e2ffc..2b0a1ea42e 100644 --- a/pyscf/tdscf/uks.py +++ b/pyscf/tdscf/uks.py @@ -45,6 +45,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: