Skip to content

Commit

Permalink
Update tdghf, tdgks, tddhf, tddks
Browse files Browse the repository at this point in the history
  • Loading branch information
@sunqm
sunqm committed Aug 27, 2024
1 parent 44902c4 commit a033b14
Show file tree
Hide file tree
Showing 12 changed files with 145 additions and 119 deletions.
32 changes: 15 additions & 17 deletions pyscf/tdscf/dhf.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -23,12 +23,11 @@
from functools import reduce
import numpy
from pyscf import lib
from pyscf import dft
from pyscf import scf
from pyscf import ao2mo
from pyscf.lib import logger
from pyscf.tdscf import rhf
from pyscf.tdscf._lr_eig import eig as lr_eig
from pyscf.data import nist
from pyscf.tdscf._lr_eig import eigh as lr_eigh, eig as lr_eig
from pyscf import __config__

OUTPUT_THRESHOLD = getattr(__config__, 'tdscf_uhf_get_nto_threshold', 0.3)
Expand Down Expand Up @@ -121,7 +120,7 @@ def add_hf_(a, b, hyb=1):
b = b - numpy.einsum('jaib->iajb', eri_mo[:nocc,nocc:,:nocc,nocc:]) * hyb
return a, b

if isinstance(mf, dft.KohnShamDFT):
if isinstance(mf, scf.hf.KohnShamDFT):
from pyscf.dft import xc_deriv
ni = mf._numint
ni.libxc.test_deriv_order(mf.xc, 2, raise_error=True)
Expand Down Expand Up @@ -399,6 +398,9 @@ def nuc_grad_method(self):

@lib.with_doc(rhf.TDA.__doc__)
class TDA(TDBase):

singlet = None

def gen_vind(self, mf=None):
'''Generate function to compute Ax'''
if mf is None:
Expand Down Expand Up @@ -450,14 +452,10 @@ def pickeig(w, v, nroots, envs):
if x0 is None:
x0 = self.init_guess(self._scf, self.nstates)

# FIXME: Is it correct to call davidson1 for complex integrals?
self.converged, self.e, x1 = \
lib.davidson1(vind, x0, precond,
tol=self.conv_tol,
nroots=nstates, lindep=self.lindep,
max_cycle=self.max_cycle,
max_space=self.max_space, pick=pickeig,
verbose=log)
self.converged, self.e, x1 = lr_eigh(
vind, x0, precond, tol_residual=self.conv_tol, lindep=self.lindep,
nroots=nstates, pick=pickeig, max_cycle=self.max_cycle,
max_memory=self.max_memory, verbose=log)

nocc = (self._scf.mo_occ>0).sum()
nmo = self._scf.mo_occ.size
Expand Down Expand Up @@ -527,7 +525,8 @@ def vind(xys):

class TDHF(TDBase):

conv_tol = 1e-5
conv_tol = 1e-4
singlet = None

@lib.with_doc(gen_tdhf_operation.__doc__)
def gen_vind(self, mf=None):
Expand Down Expand Up @@ -569,9 +568,9 @@ def pickeig(w, v, nroots, envs):
x0 = self.init_guess(self._scf, self.nstates)

self.converged, w, x1 = lr_eig(
vind, x0, precond, tol_residual=self.conv_tol, nroots=nstates,
lindep=self.lindep, max_cycle=self.max_cycle,
max_space=self.max_space, pick=pickeig, verbose=log)
vind, x0, precond, tol_residual=self.conv_tol, lindep=self.lindep,
nroots=nstates, pick=pickeig, max_cycle=self.max_cycle,
max_memory=self.max_memory, verbose=log)

nocc = (self._scf.mo_occ>0).sum()
nmo = self._scf.mo_occ.size
Expand All @@ -593,7 +592,6 @@ def norm_xy(z):
self._finalize()
return self.e, self.xy

from pyscf import scf
scf.dhf.DHF.TDA = scf.dhf.RDHF.TDA = lib.class_as_method(TDA)
scf.dhf.DHF.TDHF = scf.dhf.RDHF.TDHF = lib.class_as_method(TDHF)

Expand Down
4 changes: 0 additions & 4 deletions pyscf/tdscf/dks.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -17,12 +17,8 @@
#


import numpy
from pyscf import lib
from pyscf.tdscf import dhf
from pyscf.data import nist
from pyscf.dft.rks import KohnShamDFT
from pyscf import __config__


class TDA(dhf.TDA):
Expand Down
117 changes: 70 additions & 47 deletions pyscf/tdscf/ghf.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -25,14 +25,13 @@
from functools import reduce
import numpy
from pyscf import lib
from pyscf import dft
from pyscf.dft import numint
from pyscf import scf
from pyscf import ao2mo
from pyscf import symm
from pyscf.lib import logger
from pyscf.tdscf import rhf
from pyscf.tdscf._lr_eig import eigh as lr_eigh, eig as lr_eig
from pyscf.scf import ghf_symm
from pyscf.data import nist
from pyscf import __config__

OUTPUT_THRESHOLD = getattr(__config__, 'tdscf_rhf_get_nto_threshold', 0.3)
Expand Down Expand Up @@ -62,9 +61,7 @@ def gen_tda_operation(mf, fock_ao=None, wfnsym=None):
if isinstance(wfnsym, str):
wfnsym = symm.irrep_name2id(mol.groupname, wfnsym)
wfnsym = wfnsym % 10 # convert to D2h subgroup
orbsym = ghf_symm.get_orbsym(mol, mo_coeff)
orbsym_in_d2h = numpy.asarray(orbsym) % 10 # convert to D2h irreps
sym_forbid = (orbsym_in_d2h[occidx,None] ^ orbsym_in_d2h[viridx]) != wfnsym
sym_forbid = _get_x_sym_table(mf) != wfnsym

if fock_ao is None:
e_ia = hdiag = mo_energy[viridx] - mo_energy[occidx,None]
Expand Down Expand Up @@ -102,6 +99,14 @@ def vind(zs):
return vind, hdiag
gen_tda_hop = gen_tda_operation

def _get_x_sym_table(mf):
'''Irrep (up to D2h symmetry) of each coefficient in X[nocc,nvir]'''
mol = mf.mol
mo_occ = mf.mo_occ
orbsym = ghf_symm.get_orbsym(mol, mf.mo_coeff)
orbsym = numpy.asarray(orbsym) % 10 # convert to D2h irreps
return orbsym[mo_occ==1,None] ^ orbsym[mo_occ==0]

def get_ab(mf, mo_energy=None, mo_coeff=None, mo_occ=None):
r'''A and B matrices for TDDFT response function.
Expand Down Expand Up @@ -151,7 +156,7 @@ def add_hf_(a, b, hyb=1):
b -= numpy.einsum('jaib->iajb', eri_mo[:nocc,nocc:,:nocc,nocc:]) * hyb
return a, b

if isinstance(mf, dft.KohnShamDFT):
if isinstance(mf, scf.hf.KohnShamDFT):
from pyscf.dft import xc_deriv
ni = mf._numint
ni.libxc.test_deriv_order(mf.xc, 2, raise_error=True)
Expand Down Expand Up @@ -379,35 +384,44 @@ def gen_vind(self, mf=None):
mf = self._scf
return gen_tda_hop(mf, wfnsym=self.wfnsym)

def init_guess(self, mf, nstates=None, wfnsym=None):
def init_guess(self, mf, nstates=None, wfnsym=None, return_symmetry=False):
if nstates is None: nstates = self.nstates
if wfnsym is None: wfnsym = self.wfnsym

mo_energy = mf.mo_energy
mo_occ = mf.mo_occ
occidx = numpy.where(mo_occ==1)[0]
viridx = numpy.where(mo_occ==0)[0]
e_ia = mo_energy[viridx] - mo_energy[occidx,None]

if wfnsym is not None and mf.mol.symmetry:
if isinstance(wfnsym, str):
wfnsym = symm.irrep_name2id(mf.mol.groupname, wfnsym)
wfnsym = wfnsym % 10 # convert to D2h subgroup
orbsym = ghf_symm.get_orbsym(mf.mol, mf.mo_coeff)
orbsym_in_d2h = numpy.asarray(orbsym) % 10 # convert to D2h irreps
e_ia[(orbsym_in_d2h[occidx,None] ^ orbsym_in_d2h[viridx]) != wfnsym] = 1e99

e_ia = (mo_energy[viridx] - mo_energy[occidx,None]).ravel()
nov = e_ia.size
nstates = min(nstates, nov)
e_ia = e_ia.ravel()
e_threshold = numpy.sort(e_ia)[nstates-1]

if (wfnsym is not None or return_symmetry) and mf.mol.symmetry:
x_sym = _get_x_sym_table(mf).ravel()
if wfnsym is not None:
if isinstance(wfnsym, str):
wfnsym = symm.irrep_name2id(mf.mol.groupname, wfnsym)
wfnsym = wfnsym % 10 # convert to D2h subgroup
e_ia[x_sym != wfnsym] = 1e99
nov_allowed = numpy.count_nonzero(x_sym == wfnsym)
nstates = min(nstates, nov_allowed)

e_threshold = numpy.partition(e_ia, nstates-1)[nstates-1]
e_threshold += self.deg_eia_thresh

idx = numpy.where(e_ia <= e_threshold)[0]
x0 = numpy.zeros((idx.size, nov))
for i, j in enumerate(idx):
x0[i, j] = 1 # Koopmans' excitations
return x0

if return_symmetry:
if mf.mol.symmetry:
x0sym = x_sym[idx]
else:
x0sym = None
return x0, x0sym
else:
return x0

def kernel(self, x0=None, nstates=None):
'''TDA diagonalization solver
Expand All @@ -419,6 +433,7 @@ def kernel(self, x0=None, nstates=None):
nstates = self.nstates
else:
self.nstates = nstates
mol = self.mol

log = logger.Logger(self.stdout, self.verbose)

Expand All @@ -429,17 +444,18 @@ def pickeig(w, v, nroots, envs):
idx = numpy.where(w > self.positive_eig_threshold)[0]
return w[idx], v[:,idx], idx

x0sym = None
if x0 is None:
x0 = self.init_guess(self._scf, self.nstates)
x0, x0sym = self.init_guess(
self._scf, self.nstates, return_symmetry=True)
elif mol.symmetry:
x_sym = _get_x_sym_table(self._scf).ravel()
x0sym = [rhf._guess_wfnsym_id(self, x_sym, x) for x in x0]

# FIXME: Is it correct to call davidson1 for complex integrals
self.converged, self.e, x1 = \
lib.davidson1(vind, x0, precond,
tol=self.conv_tol,
nroots=nstates, lindep=self.lindep,
max_cycle=self.max_cycle,
max_space=self.max_space, pick=pickeig,
verbose=log)
self.converged, self.e, x1 = lr_eigh(
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 = (self._scf.mo_occ>0).sum()
nmo = self._scf.mo_occ.size
Expand Down Expand Up @@ -479,9 +495,7 @@ def gen_tdhf_operation(mf, fock_ao=None, wfnsym=None):
if isinstance(wfnsym, str):
wfnsym = symm.irrep_name2id(mol.groupname, wfnsym)
wfnsym = wfnsym % 10 # convert to D2h subgroup
orbsym = ghf_symm.get_orbsym(mol, mo_coeff)
orbsym_in_d2h = numpy.asarray(orbsym) % 10 # convert to D2h irreps
sym_forbid = (orbsym_in_d2h[occidx,None] ^ orbsym_in_d2h[viridx]) != wfnsym
sym_forbid = _get_x_sym_table(mf) != wfnsym

assert fock_ao is None

Expand Down Expand Up @@ -530,6 +544,7 @@ def vind(xys):

class TDHF(TDBase):

conv_tol = 1e-4
singlet = None

@lib.with_doc(gen_tdhf_operation.__doc__)
Expand All @@ -538,10 +553,15 @@ def gen_vind(self, mf=None):
mf = self._scf
return gen_tdhf_operation(mf, wfnsym=self.wfnsym)

def init_guess(self, mf, nstates=None, wfnsym=None):
x0 = TDA.init_guess(self, mf, nstates, wfnsym)
y0 = numpy.zeros_like(x0)
return numpy.asarray(numpy.block([[x0, y0], [y0, x0.conj()]]))
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)
return numpy.hstack([x0, y0]), x0sym
else:
x0 = TDA.init_guess(self, mf, nstates, wfnsym, return_symmetry)
y0 = numpy.zeros_like(x0)
return numpy.hstack([x0, y0])

def kernel(self, x0=None, nstates=None):
'''TDHF diagonalization with non-Hermitian eigenvalue solver
Expand All @@ -553,6 +573,7 @@ def kernel(self, x0=None, nstates=None):
nstates = self.nstates
else:
self.nstates = nstates
mol = self.mol

log = logger.Logger(self.stdout, self.verbose)

Expand All @@ -566,16 +587,19 @@ def pickeig(w, v, nroots, envs):
# FIXME: Should the amplitudes be real? It also affects x2c-tdscf
return lib.linalg_helper._eigs_cmplx2real(w, v, realidx, ensure_real)

x0sym = None
if x0 is None:
x0 = self.init_guess(self._scf, self.nstates)

self.converged, w, x1 = \
lib.davidson_nosym1(vind, x0, precond,
tol=self.conv_tol,
nroots=nstates, lindep=self.lindep,
max_cycle=self.max_cycle,
max_space=self.max_space, pick=pickeig,
verbose=log)
x0, x0sym = self.init_guess(
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 = [rhf._guess_wfnsym_id(self, x_sym, x) for x in x0]

self.converged, w, 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)

nocc = (self._scf.mo_occ>0).sum()
nmo = self._scf.mo_occ.size
Expand All @@ -598,6 +622,5 @@ def norm_xy(z):

RPA = TDGHF = TDHF

from pyscf import scf
scf.ghf.GHF.TDA = lib.class_as_method(TDA)
scf.ghf.GHF.TDHF = lib.class_as_method(TDHF)
32 changes: 16 additions & 16 deletions pyscf/tdscf/gks.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -20,10 +20,8 @@
import numpy
from pyscf import lib
from pyscf import symm
from pyscf.tdscf import ghf
from pyscf.scf import ghf_symm
from pyscf.scf import _response_functions # noqa
from pyscf.data import nist
from pyscf.tdscf import ghf, rhf
from pyscf.tdscf._lr_eig import eigh as lr_eigh
from pyscf.dft.rks import KohnShamDFT
from pyscf import __config__

Expand All @@ -39,6 +37,7 @@ class TDDFT(ghf.TDHF):
class CasidaTDDFT(TDDFT, TDA):
'''Solve the Casida TDDFT formula (A-B)(A+B)(X+Y) = (X+Y)w^2
'''

init_guess = TDA.init_guess

def gen_vind(self, mf=None):
Expand All @@ -62,9 +61,7 @@ def gen_vind(self, mf=None):
if isinstance(wfnsym, str):
wfnsym = symm.irrep_name2id(mol.groupname, wfnsym)
wfnsym = wfnsym % 10 # convert to D2h subgroup
orbsym = ghf_symm.get_orbsym(mol, mo_coeff)
orbsym_in_d2h = numpy.asarray(orbsym) % 10 # convert to D2h irreps
sym_forbid = (orbsym_in_d2h[occidx,None] ^ orbsym_in_d2h[viridx]) != wfnsym
sym_forbid = ghf._get_x_sym_table(mf) != wfnsym

e_ia = (mo_energy[viridx].reshape(-1,1) - mo_energy[occidx]).T
if wfnsym is not None and mol.symmetry:
Expand Down Expand Up @@ -101,6 +98,7 @@ def kernel(self, x0=None, nstates=None):
'''TDDFT diagonalization solver
'''
cpu0 = (lib.logger.process_clock(), lib.logger.perf_counter())
mol = self.mol
mf = self._scf
if mf._numint.libxc.is_hybrid_xc(mf.xc):
raise RuntimeError('%s cannot be used with hybrid functional'
Expand All @@ -121,16 +119,18 @@ def pickeig(w, v, nroots, envs):
idx = numpy.where(w > self.positive_eig_threshold)[0]
return w[idx], v[:,idx], idx

x0sym = None
if x0 is None:
x0 = self.init_guess(self._scf, self.nstates)

self.converged, w2, x1 = \
lib.davidson1(vind, x0, precond,
tol=self.conv_tol,
nroots=nstates, lindep=self.lindep,
max_cycle=self.max_cycle,
max_space=self.max_space, pick=pickeig,
verbose=log)
x0, x0sym = self.init_guess(
self._scf, self.nstates, return_symmetry=True)
elif mol.symmetry:
x_sym = ghf._get_x_sym_table(self._scf).ravel()
x0sym = [rhf._guess_wfnsym_id(self, x_sym, x) for x in x0]

self.converged, w2, x1 = lr_eigh(
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)

mo_energy = self._scf.mo_energy
mo_occ = self._scf.mo_occ
Expand Down
Loading

0 comments on commit a033b14

Please sign in to comment.