Optimize nuclear hessian (pyscf#2241)

* Remove chkfile dependency in hessian

* Update krylov solver

* Update cphf

* Update hessian code

* Update eph hessian

* Fix a shape bug in ucphf

* Remove duplicated code

* clean up

* Adjust tests

* Update docstring of lib.linalg_helper.solve

examples/eph/00-simple_eph.py

@@ -3,7 +3,31 @@
 '''
 A simple example to run EPH calculation.
 '''
+
 from pyscf import gto, dft, eph
+
+mol = gto.M()
+mol.atom = [['O', [0.000000000000,  -0.000000000775,   0.923671924285]],
+            ['H', [-0.000000000000,  -1.432564848017,   2.125164039823]],
+            ['H', [0.000000000000,   1.432564848792,   2.125164035930]]]
+mol.unit = 'Bohr'
+mol.basis = 'sto3g'
+mol.verbose = 4
+mol.build() # this is a pre-computed relaxed geometry
+
+mf = mol.RHF()
+mf.conv_tol = 1e-16
+mf.conv_tol_grad = 1e-10
+mf.kernel()
+
+myeph = eph.EPH(mf)
+grad = mf.nuc_grad_method().kernel()
+print("Force on the atoms/au:")
+print(grad)
+eph, omega = myeph.kernel(mo_rep=True)
+print(np.amax(eph))
+
+
 mol = gto.M(atom='N 0 0 0; N 0 0 2.100825', basis='def2-svp', verbose=4, unit="bohr")
 # this is a pre-computed relaxed molecule
 # for geometry relaxation, refer to pyscf/example/geomopt

pyscf/df/hessian/rhf.py

@@ -368,16 +368,8 @@ def make_h1(hessobj, mo_coeff, mo_occ, chkfile=None, atmlst=None, verbose=None):
     for ia, h1, vj1, vk1 in _gen_jk(hessobj, mo_coeff, mo_occ, chkfile,
                                     atmlst, verbose, True):
         h1 += vj1 - vk1 * .5
-
-        if chkfile is None:
-            h1ao[ia] = h1
-        else:
-            key = 'scf_f1ao/%d' % ia
-            lib.chkfile.save(chkfile, key, h1)
-    if chkfile is None:
-        return h1ao
-    else:
-        return chkfile
+        h1ao[ia] = h1
+    return h1ao
 
 def _gen_jk(hessobj, mo_coeff, mo_occ, chkfile=None, atmlst=None,
             verbose=None, with_k=True):

pyscf/df/hessian/rks.py

@@ -109,13 +109,7 @@ def make_h1(hessobj, mo_coeff, mo_occ, chkfile=None, atmlst=None, verbose=None):
             for ia, h1, vj1, vk1 in df_rhf_hess._gen_jk(
                     hessobj, mo_coeff, mo_occ, chkfile, atmlst, verbose):
                 h1ao[ia] -= .5 * (alpha - hyb) * vk1
-
-    if chkfile is None:
-        return h1ao
-    else:
-        for ia in atmlst:
-            lib.chkfile.save(chkfile, 'scf_f1ao/%d'%ia, h1ao[ia])
-        return chkfile
+    return h1ao
 
 
 class Hessian(rks_hess.Hessian):

pyscf/df/hessian/uhf.py

@@ -398,19 +398,9 @@ def make_h1(hessobj, mo_coeff, mo_occ, chkfile=None, atmlst=None, verbose=None):
     for ia, h1, vj1, vk1 in _gen_jk(hessobj, mo_coeff, mo_occ, chkfile,
                                     atmlst, verbose, True):
         vk1a, vk1b = vk1
-        h1a = h1 + vj1 - vk1a
-        h1b = h1 + vj1 - vk1b
-
-        if chkfile is None:
-            h1aoa[ia] = h1a
-            h1aob[ia] = h1b
-        else:
-            lib.chkfile.save(chkfile, 'scf_f1ao/0/%d' % ia, h1a)
-            lib.chkfile.save(chkfile, 'scf_f1ao/1/%d' % ia, h1b)
-    if chkfile is None:
-        return (h1aoa,h1aob)
-    else:
-        return chkfile
+        h1aoa[ia] = h1 + vj1 - vk1a
+        h1aob[ia] = h1 + vj1 - vk1b
+    return (h1aoa,h1aob)
 
 def _gen_jk(hessobj, mo_coeff, mo_occ, chkfile=None, atmlst=None,
             verbose=None, with_k=True):

pyscf/df/hessian/uks.py

@@ -121,14 +121,7 @@ def make_h1(hessobj, mo_coeff, mo_occ, chkfile=None, atmlst=None, verbose=None):
                 vk1a, vk1b = vk1
                 h1aoa[ia] -= (alpha - hyb) * vk1a
                 h1aob[ia] -= (alpha - hyb) * vk1b
-
-    if chkfile is None:
-        return h1aoa, h1aob
-    else:
-        for ia in atmlst:
-            lib.chkfile.save(chkfile, 'scf_f1ao/0/%d'%ia, h1aoa[ia])
-            lib.chkfile.save(chkfile, 'scf_f1ao/1/%d'%ia, h1aob[ia])
-        return chkfile
+    return h1aoa, h1aob
 
 
 class Hessian(uks_hess.Hessian):

pyscf/eph/rhf.py

@@ -23,7 +23,7 @@
 import numpy as np
 import scipy.linalg
 from pyscf.hessian import rhf
-from pyscf.lib import logger, chkfile
+from pyscf.lib import logger
 from pyscf.scf._response_functions import _gen_rhf_response
 from pyscf import __config__
 from pyscf.data.nist import HARTREE2WAVENUMBER, MP_ME
@@ -37,16 +37,16 @@ def kernel(ephobj, mo_energy=None, mo_coeff=None, mo_occ=None, mo_rep=False):
     if mo_coeff is None: mo_coeff = ephobj.base.mo_coeff
     if mo_occ is None: mo_occ = ephobj.base.mo_occ
 
-    # chkfile is used to pass first orbitals from hessian methods to eph methods
-    # TODO: Remove the dependence to chfile and return first orbitals from a function
-    assert ephobj.chkfile is not None, 'chkfile is required to save first order orbitals'
+    h1ao = ephobj.make_h1(mo_coeff, mo_occ)
+    mo1, mo_e1 = ephobj.solve_mo1(mo_energy, mo_coeff, mo_occ, h1ao)
 
-    de = ephobj.hess_elec(mo_energy, mo_coeff, mo_occ)
+    de = ephobj.hess_elec(mo_energy, mo_coeff, mo_occ,
+                          mo1=mo1, mo_e1=mo_e1, h1ao=h1ao)
     ephobj.de = de + ephobj.hess_nuc(ephobj.mol)
 
     omega, vec = ephobj.get_mode(ephobj.mol, ephobj.de)
     ephobj.omega, ephobj.vec = omega, vec
-    ephobj.eph = ephobj.get_eph(ephobj.chkfile, omega, vec, mo_rep)
+    ephobj.eph = ephobj.get_eph(mo1, omega, vec, mo_rep)
     return ephobj.eph, ephobj.omega
 
 def solve_hmat(mol, hmat, cutoff_frequency=CUTOFF_FREQUENCY,
@@ -143,10 +143,6 @@ def _freq_mass_weighted_vec(vec, omega, mass):
     return vec
 
 def get_eph(ephobj, mo1, omega, vec, mo_rep):
-    if isinstance(mo1, str):
-        mo1 = chkfile.load(mo1, 'scf_mo1')
-        mo1 = {int(k): mo1[k] for k in mo1}
-
     mol = ephobj.mol
     mf = ephobj.base
     vnuc_deriv = ephobj.vnuc_generator(mol)
@@ -209,29 +205,3 @@ def __init__(self, scf_method, cutoff_frequency=CUTOFF_FREQUENCY,
     get_eph = get_eph
     vnuc_generator = vnuc_generator
     kernel = kernel
-
-if __name__ == '__main__':
-    from pyscf import gto, scf
-
-    mol = gto.M()
-    mol.atom = [['O', [0.000000000000,  -0.000000000775,   0.923671924285]],
-                ['H', [-0.000000000000,  -1.432564848017,   2.125164039823]],
-                ['H', [0.000000000000,   1.432564848792,   2.125164035930]]]
-    mol.unit = 'Bohr'
-    mol.basis = 'sto3g'
-    mol.verbose=4
-    mol.build() # this is a pre-computed relaxed geometry
-
-    mf = scf.RHF(mol)
-    mf.conv_tol = 1e-16
-    mf.conv_tol_grad = 1e-10
-    mf.kernel()
-
-    myeph = EPH(mf)
-
-    grad = mf.nuc_grad_method().kernel()
-    print("Force on the atoms/au:")
-    print(grad)
-
-    eph, omega = myeph.kernel(mo_rep=True)
-    print(np.amax(eph))

pyscf/eph/rks.py

@@ -100,10 +100,6 @@ def _get_vxc_deriv1(hessobj, mo_coeff, mo_occ, max_memory):
     return vmat
 
 def get_eph(ephobj, mo1, omega, vec, mo_rep):
-    if isinstance(mo1, str):
-        mo1 = lib.chkfile.load(mo1, 'scf_mo1')
-        mo1 = {int(k): mo1[k] for k in mo1}
-
     mol = ephobj.mol
     mf = ephobj.base
     ni = mf._numint

pyscf/eph/uhf.py

@@ -54,18 +54,12 @@ def fx(mo1):
     return fx
 
 def get_eph(ephobj, mo1, omega, vec, mo_rep):
-    if isinstance(mo1, str):
-        mo1 = lib.chkfile.load(mo1, 'scf_mo1')
-        mo1a = mo1['0']
-        mo1b = mo1['1']
-        mo1a = {int(k): mo1a[k] for k in mo1a}
-        mo1b = {int(k): mo1b[k] for k in mo1b}
-
     mol = ephobj.mol
     mf = ephobj.base
     vnuc_deriv = ephobj.vnuc_generator(mol)
     aoslices = mol.aoslice_by_atom()
 
+    mo1a, mo1b = mo1
     mo_coeff, mo_occ = mf.mo_coeff, mf.mo_occ
     vind = uhf_deriv_generator(mf, mf.mo_coeff, mf.mo_occ)
     nao, nmo = mo_coeff[0].shape

pyscf/eph/uks.py

@@ -124,13 +124,6 @@ def _get_vxc_deriv1(hessobj, mo_coeff, mo_occ, max_memory):
     return vmata, vmatb
 
 def get_eph(ephobj, mo1, omega, vec, mo_rep):
-    if isinstance(mo1, str):
-        mo1 = lib.chkfile.load(mo1, 'scf_mo1')
-        mo1a = mo1['0']
-        mo1b = mo1['1']
-        mo1a = {int(k): mo1a[k] for k in mo1a}
-        mo1b = {int(k): mo1b[k] for k in mo1b}
-
     mol = ephobj.mol
     mf = ephobj.base
     ni = mf._numint
@@ -141,6 +134,7 @@ def get_eph(ephobj, mo1, omega, vec, mo_rep):
     vnuc_deriv = ephobj.vnuc_generator(mol)
     aoslices = mol.aoslice_by_atom()
 
+    mo1a, mo1b = mo1
     mo_coeff, mo_occ = mf.mo_coeff, mf.mo_occ
     vind = uhf_deriv_generator(mf, mf.mo_coeff, mf.mo_occ)
     mem_now = lib.current_memory()[0]