Add level shift to CPHF solver

pyscf/dft/xcfun.py

@@ -252,7 +252,7 @@
           'B97XC'   , 'B97_1XC' , 'B97_2XC' , 'M05XC'   , 'TPSSH'   ,
           'HFLYP'}
 RSH_XC = {'CAMB3LYP'}
-MAX_DERIV_ORDER = ctypes.c_int.in_dll(_itrf, 'XCFUN_max_deriv_order').value
+MAX_DERIV_ORDER = 3
 
 VV10_XC = {
     'B97M_V'    : (6.0, 0.01),
@@ -338,8 +338,12 @@ def rsh_coeff(xc_code):
     beta = hyb - alpha
     return omega, alpha, beta
 
+@lru_cache(100)
 def max_deriv_order(xc_code):
-    return MAX_DERIV_ORDER
+    try:
+        return ctypes.c_int.in_dll(_itrf, 'XCFUN_max_deriv_order').value
+    except ValueError: # make it compatible with old xcfun_itrf.so
+        return MAX_DERIV_ORDER
 
 def test_deriv_order(xc_code, deriv, raise_error=False):
     support = deriv <= max_deriv_order(xc_code)

pyscf/hessian/rhf.py

@@ -296,7 +296,8 @@ def _get_jk(mol, intor, comp, aosym, script_dms,
     return vs
 
 def solve_mo1(mf, mo_energy, mo_coeff, mo_occ, h1ao_or_chkfile,
-              fx=None, atmlst=None, max_memory=4000, verbose=None, max_cycle=50):
+              fx=None, atmlst=None, max_memory=4000, verbose=None,
+              max_cycle=50, level_shift=0):
     '''Solve the first order equation
 
     Kwargs:
@@ -343,7 +344,8 @@ def _ao2mo(mat):
 
         h1vo = numpy.vstack(h1vo)
         s1vo = numpy.vstack(s1vo)
-        mo1, e1 = cphf.solve(fx, mo_energy, mo_occ, h1vo, s1vo, max_cycle=max_cycle)
+        mo1, e1 = cphf.solve(fx, mo_energy, mo_occ, h1vo, s1vo,
+                             max_cycle=max_cycle, level_shift=level_shift)
         mo1 = numpy.einsum('pq,xqi->xpi', mo_coeff, mo1).reshape(-1,3,nao,nocc)
         e1 = e1.reshape(-1,3,nocc,nocc)
 
@@ -470,8 +472,15 @@ def h_op(x):
 class HessianBase(lib.StreamObject):
     '''Non-relativistic restricted Hartree-Fock hessian'''
 
+    # Max. number of iterations for Krylov solver
+    max_cycle = 50
+    # Shift virtual orbitals to slightly improve the convergence speed of Krylov solver
+    # A small level_shift ~ 0.1 is often helpful to decrease 2 - 3 iterations
+    # while the error of cphf solver may be increased by one magnitude.
+    level_shift = 0
+
     _keys = {
-        'mol', 'base', 'chkfile', 'atmlst', 'de', 'max_cycle'
+        'mol', 'base', 'chkfile', 'atmlst', 'de', 'max_cycle', 'level_shift'
     }
 
     def __init__(self, scf_method):
@@ -481,7 +490,6 @@ def __init__(self, scf_method):
         self.base = scf_method
         self.chkfile = scf_method.chkfile
         self.max_memory = self.mol.max_memory
-        self.max_cycle = 50
         self.atmlst = range(self.mol.natm)
         self.de = numpy.zeros((0,0,3,3))  # (A,B,dR_A,dR_B)
 
@@ -566,7 +574,8 @@ def get_hcore(iatm, jatm):
     def solve_mo1(self, mo_energy, mo_coeff, mo_occ, h1ao_or_chkfile,
                   fx=None, atmlst=None, max_memory=4000, verbose=None):
         return solve_mo1(self.base, mo_energy, mo_coeff, mo_occ, h1ao_or_chkfile,
-                         fx, atmlst, max_memory, verbose, max_cycle=self.max_cycle)
+                         fx, atmlst, max_memory, verbose,
+                         max_cycle=self.max_cycle, level_shift=self.level_shift)
 
     def hess_nuc(self, mol=None, atmlst=None):
         if mol is None: mol = self.mol

pyscf/hessian/test/test_rhf.py

@@ -35,6 +35,18 @@ def tearDownModule():
     del mol
 
 class KnownValues(unittest.TestCase):
+    def test_rhf_hess(self):
+        mf = scf.RHF(mol)
+        e0 = mf.kernel()
+        hess = hessian.RHF(mf).kernel()
+        self.assertAlmostEqual(lib.fp(hess), -0.7816352153153946, 6)
+
+        hobj = hessian.RHF(mf)
+        hobj.max_cycle = 10
+        hobj.level_shift = .1
+        hess = hobj.kernel()
+        self.assertAlmostEqual(lib.fp(hess), -0.7816352153153946, 6)
+
     def test_finite_diff_x2c_rhf_hess(self):
         mf = scf.RHF(mol).x2c()
         mf.conv_tol = 1e-14

pyscf/hessian/test/test_uhf.py

@@ -36,6 +36,15 @@ def tearDownModule():
     del mol
 
 class KnownValues(unittest.TestCase):
+    def test_uhf_hess(self):
+        mf = scf.UHF(mol)
+        mf.conv_tol = 1e-14
+        e0 = mf.kernel()
+        hobj = mf.Hessian()
+        hobj.level_shift = .05
+        hess = hobj.kernel()
+        self.assertAlmostEqual(lib.fp(hess), -0.20243405976628576, 5)
+
     def test_finite_diff_rhf_hess(self):
         mf = scf.UHF(mol)
         mf.conv_tol = 1e-14

pyscf/hessian/uhf.py

@@ -257,7 +257,8 @@ def make_h1(hessobj, mo_coeff, mo_occ, chkfile=None, atmlst=None, verbose=None):
         return chkfile
 
 def solve_mo1(mf, mo_energy, mo_coeff, mo_occ, h1ao_or_chkfile,
-              fx=None, atmlst=None, max_memory=4000, verbose=None):
+              fx=None, atmlst=None, max_memory=4000, verbose=None,
+              max_cycle=50, level_shift=0):
     mol = mf.mol
     if atmlst is None: atmlst = range(mol.natm)
 
@@ -306,7 +307,8 @@ def _ao2mo(mat, mo_coeff, mocc):
 
         h1vo = (numpy.vstack(h1voa), numpy.vstack(h1vob))
         s1vo = (numpy.vstack(s1voa), numpy.vstack(s1vob))
-        mo1, e1 = ucphf.solve(fx, mo_energy, mo_occ, h1vo, s1vo)
+        mo1, e1 = ucphf.solve(fx, mo_energy, mo_occ, h1vo, s1vo,
+                              max_cycle=max_cycle, level_shift=level_shift)
         mo1a = numpy.einsum('pq,xqi->xpi', mo_coeff[0], mo1[0]).reshape(-1,3,nao,nocca)
         mo1b = numpy.einsum('pq,xqi->xpi', mo_coeff[1], mo1[1]).reshape(-1,3,nao,noccb)
         e1a = e1[0].reshape(-1,3,nocca,nocca)
@@ -449,7 +451,8 @@ class Hessian(rhf_hess.HessianBase):
     def solve_mo1(self, mo_energy, mo_coeff, mo_occ, h1ao_or_chkfile,
                   fx=None, atmlst=None, max_memory=4000, verbose=None):
         return solve_mo1(self.base, mo_energy, mo_coeff, mo_occ, h1ao_or_chkfile,
-                         fx, atmlst, max_memory, verbose)
+                         fx, atmlst, max_memory, verbose,
+                         max_cycle=self.max_cycle, level_shift=self.level_shift)
 
     def to_gpu(self):
         raise NotImplementedError

pyscf/lib/linalg_helper.py

@@ -1290,9 +1290,9 @@ def krylov(aop, b, x0=None, tol=1e-10, max_cycle=30, dot=numpy.dot,
     >>> from pyscf import lib
     >>> a = numpy.random.random((10,10)) * 1e-2
     >>> b = numpy.random.random(10)
-    >>> aop = lambda x: numpy.dot(a,x)
+    >>> aop = lambda x: a.dot(x.T).T
     >>> x = lib.krylov(aop, b)
-    >>> numpy.allclose(numpy.dot(a,x)+x, b)
+    >>> numpy.allclose(aop(x)+x, b)
     True
     '''
     if isinstance(aop, numpy.ndarray) and aop.ndim == 2:

pyscf/lib/test/test_linalg_helper.py

@@ -116,6 +116,24 @@ def precond(x, *args):
         x1 = linalg_helper.krylov(aop, b/a_diag, x1, max_cycle=30)
         self.assertAlmostEqual(abs(xref - x1).max(), 0, 6)
 
+    def test_krylov_with_level_shift(self):
+        numpy.random.seed(10)
+        n = 100
+        a = numpy.random.rand(n,n) * .1
+        a = a.dot(a.T)
+        a_diag = numpy.random.rand(n)
+        b = numpy.random.rand(n)
+        ref = numpy.linalg.solve(numpy.diag(a_diag) + a, b)
+
+        #((diag+shift) + (a-shift)) x = b
+        shift = .1
+        a_diag += shift
+        a -= numpy.eye(n)*shift
+
+        aop = lambda x: (a.dot(x.T).T/a_diag)
+        c = linalg_helper.krylov(aop, b/a_diag, max_cycle=18)
+        self.assertAlmostEqual(abs(ref - c).max(), 0, 9)
+
     def test_dgeev(self):
         numpy.random.seed(12)
         n = 100

pyscf/scf/cphf.py

@@ -27,7 +27,8 @@
 
 
 def solve(fvind, mo_energy, mo_occ, h1, s1=None,
-          max_cycle=20, tol=1e-9, hermi=False, verbose=logger.WARN):
+          max_cycle=50, tol=1e-9, hermi=False, verbose=logger.WARN,
+          level_shift=0):
     '''
     Args:
         fvind : function
@@ -36,29 +37,43 @@ def solve(fvind, mo_energy, mo_occ, h1, s1=None,
     Kwargs:
         hermi : boolean
             Whether the matrix defined by fvind is Hermitian or not.
+        level_shift : float
+            Add to diagonal terms to slightly improve the convergence speed of
+            Krylov solver
     '''
     if s1 is None:
         return solve_nos1(fvind, mo_energy, mo_occ, h1,
-                          max_cycle, tol, hermi, verbose)
+                          max_cycle, tol, hermi, verbose, level_shift)
     else:
         return solve_withs1(fvind, mo_energy, mo_occ, h1, s1,
-                            max_cycle, tol, hermi, verbose)
+                            max_cycle, tol, hermi, verbose, level_shift)
 kernel = solve
 
 # h1 shape is (:,nvir,nocc)
 def solve_nos1(fvind, mo_energy, mo_occ, h1,
-               max_cycle=20, tol=1e-9, hermi=False, verbose=logger.WARN):
-    '''For field independent basis. First order overlap matrix is zero'''
+               max_cycle=50, tol=1e-9, hermi=False, verbose=logger.WARN,
+               level_shift=0):
+    '''For field independent basis. First order overlap matrix is zero
+
+    Kwargs:
+        level_shift : float
+            Add to diagonal terms to slightly improve the convergence speed of
+            Krylov solver
+    '''
+    assert not hermi
     log = logger.new_logger(verbose=verbose)
     t0 = (logger.process_clock(), logger.perf_counter())
 
     e_a = mo_energy[mo_occ==0]
     e_i = mo_energy[mo_occ>0]
-    e_ai = 1 / lib.direct_sum('a-i->ai', e_a, e_i)
+    e_ai = 1 / (e_a[:,None] + level_shift - e_i)
     mo1base = h1 * -e_ai
 
     def vind_vo(mo1):
-        v = fvind(mo1.reshape(h1.shape)).reshape(h1.shape)
+        mo1 = mo1.reshape(h1.shape)
+        v = fvind(mo1).reshape(h1.shape)
+        if level_shift != 0:
+            v -= mo1 * level_shift
         v *= e_ai
         return v.ravel()
     mo1 = lib.krylov(vind_vo, mo1base.ravel(),
@@ -68,55 +83,62 @@ def vind_vo(mo1):
 
 # h1 shape is (:,nocc+nvir,nocc)
 def solve_withs1(fvind, mo_energy, mo_occ, h1, s1,
-                 max_cycle=20, tol=1e-9, hermi=False, verbose=logger.WARN):
+                 max_cycle=50, tol=1e-9, hermi=False, verbose=logger.WARN,
+                 level_shift=0):
     '''For field dependent basis. First order overlap matrix is non-zero.
     The first order orbitals are set to
     C^1_{ij} = -1/2 S1
     e1 = h1 - s1*e0 + (e0_j-e0_i)*c1 + vhf[c1]
 
     Kwargs:
-        hermi : boolean
-            Whether the matrix defined by fvind is Hermitian or not.
+        level_shift : float
+            Add to diagonal terms to slightly improve the convergence speed of
+            Krylov solver
 
     Returns:
         First order orbital coefficients (in MO basis) and first order orbital
         energy matrix
     '''
+    assert not hermi
     log = logger.new_logger(verbose=verbose)
     t0 = (logger.process_clock(), logger.perf_counter())
 
     occidx = mo_occ > 0
     viridx = mo_occ == 0
     e_a = mo_energy[viridx]
     e_i = mo_energy[occidx]
-    e_ai = 1 / lib.direct_sum('a-i->ai', e_a, e_i)
+    e_ai = 1 / (e_a[:,None] + level_shift - e_i)
     nvir, nocc = e_ai.shape
     nmo = nocc + nvir
 
     s1 = s1.reshape(-1,nmo,nocc)
     hs = mo1base = h1.reshape(-1,nmo,nocc) - s1*e_i
-    mo_e1 = hs[:,occidx,:].copy()
 
+    mo1base = hs.copy()
     mo1base[:,viridx] *= -e_ai
     mo1base[:,occidx] = -s1[:,occidx] * .5
 
     def vind_vo(mo1):
-        v = fvind(mo1.reshape(h1.shape)).reshape(-1,nmo,nocc)
+        mo1 = mo1.reshape(mo1base.shape)
+        v = fvind(mo1).reshape(mo1base.shape)
+        if level_shift != 0:
+            v -= mo1 * level_shift
         v[:,viridx,:] *= e_ai
         v[:,occidx,:] = 0
         return v.ravel()
     mo1 = lib.krylov(vind_vo, mo1base.ravel(),
                      tol=tol, max_cycle=max_cycle, hermi=hermi, verbose=log)
     mo1 = mo1.reshape(mo1base.shape)
+    mo1[:,occidx] = mo1base[:,occidx]
     log.timer('krylov solver in CPHF', *t0)
 
-    v1mo = fvind(mo1.reshape(h1.shape)).reshape(-1,nmo,nocc)
-    mo1[:,viridx] = mo1base[:,viridx] - v1mo[:,viridx]*e_ai
+    hs += fvind(mo1).reshape(mo1base.shape)
+    mo1[:,viridx] = hs[:,viridx] / (e_i - e_a[:,None])
 
     # mo_e1 has the same symmetry as the first order Fock matrix (hermitian or
     # anti-hermitian). mo_e1 = v1mo - s1*lib.direct_sum('i+j->ij',e_i,e_i)
-    mo_e1 += mo1[:,occidx] * lib.direct_sum('i-j->ij', e_i, e_i)
-    mo_e1 += v1mo[:,occidx,:]
+    mo_e1 = hs[:,occidx,:]
+    mo_e1 += mo1[:,occidx] * (e_i[:,None] - e_i)
 
     if h1.ndim == 3:
         return mo1, mo_e1