Feat: Add trust-region QN

molbe/_opt.py

@@ -118,7 +118,7 @@ def objfunc(self, xk):
         return errvec_
 
 
-    def optimize(self, method, J0 = None):
+    def optimize(self, method, J0 = None, trust_region=False):
         """Main kernel to perform BE optimization
 
         Parameters
@@ -127,6 +127,8 @@ def optimize(self, method, J0 = None):
            High-level quantum chemistry method.
         J0 : list of list of float, optional
            Initial Jacobian
+        trust_region : bool, optional
+           Use trust-region based QN optimization, by default False
         """
         from molbe.external.optqn import FrankQN
         import sys
@@ -155,19 +157,24 @@ def optimize(self, method, J0 = None):
             optQN = FrankQN(self.objfunc, numpy.array(self.pot),
                             f0, J0,
                             max_space=self.max_space)
-
-            # Perform optimization steps
-            for iter_ in range(self.max_space):
-                optQN.next_step()
-                self.iter += 1
-                print('-- In iter ',self.iter, flush=True)    
-                print('Error in density matching      :   {:>2.4e}'.format(self.err), flush=True)
+
+            if self.err < self.conv_tol:
                 print(flush=True)
-                if self.err < self.conv_tol:
-                    print(flush=True)
-                    print('CONVERGED',flush=True)
+                print('CONVERGED w/o Optimization Steps',flush=True)
+                print(flush=True)
+            else:
+                # Perform optimization steps
+                for iter_ in range(self.max_space):
+                    print('-- In iter ',self.iter, flush=True)
+                    optQN.next_step(trust_region=trust_region)
+                    self.iter += 1
+                    print('Error in density matching      :   {:>2.4e}'.format(self.err), flush=True)
                     print(flush=True)
-                    break
+                    if self.err < self.conv_tol:
+                        print(flush=True)
+                        print('CONVERGED',flush=True)
+                        print(flush=True)
+                        break
         else:
             print('This optimization method for BE is not supported')
             sys.exit()
@@ -176,8 +183,9 @@ def optimize(self, method, J0 = None):
 
 
 def optimize(self, solver='MP2',method='QN',
-             only_chem=False, conv_tol = 1.e-6,relax_density=False, use_cumulant=True,
-             J0=None, nproc=1, ompnum=4, max_iter=500, scratch_dir=None,  **solver_kwargs):
+             only_chem=False, conv_tol = 1.e-6, relax_density=False, use_cumulant=True,
+             J0=None, nproc=1, ompnum=4, max_iter=500, scratch_dir=None, trust_region=False,
+             **solver_kwargs):
     """BE optimization function
 
     Interfaces BEOPT to perform bootstrap embedding optimization.
@@ -207,6 +215,8 @@ def optimize(self, solver='MP2',method='QN',
        If nproc > 1, ompnum sets the number of cores for OpenMP parallelization. Defaults to 4
     J0 : list of list of float
        Initial Jacobian.
+    trust_region : bool, optional
+       Use trust-region based QN optimization, by default False
     """
     from .misc import print_energy
 
@@ -239,7 +249,7 @@ def optimize(self, solver='MP2',method='QN',
             J0 = self.get_be_error_jacobian(jac_solver='HF')                        
 
         # Perform the optimization
-        be_.optimize(method, J0=J0)
+        be_.optimize(method, J0=J0, trust_region=trust_region)
         self.ebe_tot = self.ebe_hf + be_.Ebe[0]
         # Print the energy components
         print_energy(be_.Ebe[0], be_.Ebe[1][1], be_.Ebe[1][0]+be_.Ebe[1][2], self.ebe_hf)

molbe/external/optqn.py

@@ -1,10 +1,12 @@
 # Author(s): Hong-Zhou Ye
 #            Oinam Romesh Meitei
-# NOTICE: The following code is entirely written by Hong-Zhou Ye.
+#            Minsik Cho
+# NOTICE: The following code is mostly written by Hong-Zhou Ye (except for the trust region routine)
 #         The code has been slightly modified.
 #
 
 import numpy,sys, h5py
+from .. import be_var
 
 
 def line_search_LF(func, xold, fold, dx, iter_):
@@ -45,10 +47,76 @@ def line_search_LF(func, xold, fold, dx, iter_):
     print(flush=True)
     return alp, xk, fk
 
-def opt_QN(self):
-
-    J0 = self.get_be_error_jacobian(solver)
-
+def trustRegion(func, xold, fold, Binv, c = 0.5):
+    """Perform Trust Region Optimization. See "A Broyden Trust Region Quasi-Newton Method 
+    for Nonlinear Equations" (https://www.iaeng.org/IJCS/issues_v46/issue_3/IJCS_46_3_09.pdf)
+    Algorithm 1 for more information
+
+    Parameters
+    ----------
+    func : function
+        Cost function
+    xold : list or numpy.ndarray
+        Current x_p (potentials in BE optimization)
+    fold : list or numpy.ndarray
+        Current f(x_p) (error vector)
+    Binv : numpy.ndarray
+        Inverse of Jacobian approximate (B^{-1}); This is updated in Broyden's Method through Sherman-Morrison formula
+    c : float, optional
+        Initial value of trust radius ∈ (0, 1), by default 0.5
+
+    Returns
+    -------
+    xnew, fnew
+        x_{p+1} and f_{p+1}. These values are used to proceed with Broyden's Method.
+    """
+    # c := initial trust radius (trust_radius = c^p)
+    microiter = 0 # p
+    rho = 0.001 # Threshold for trust region subproblem
+    ratio = 0 # Initial r
+    B = numpy.linalg.inv(Binv) # approx Jacobian
+    #dx_gn = - Binv@fold
+    dx_gn = -(Binv@Binv.T)@B.T@fold
+    dx_sd = - B.T@fold # Steepest Descent step
+    t = numpy.linalg.norm(dx_sd)**2 / numpy.linalg.norm(B@dx_sd)**2
+    prevdx = None
+    while (ratio < rho or ared < 0.):
+        # Trust Region subproblem
+        # minimize (1/2) ||F_k + B_k d||^2 w.r.t. d, s.t. d w/i trust radius
+        # to pick the optimal direction using dog leg method
+        if numpy.linalg.norm(dx_gn) < max(1.0, numpy.linalg.norm(xold)) * (c ** microiter): # Gauss-Newton step within the trust radius
+            print('  Trust Region Optimization Step ', microiter, ': Gauss-Newton', flush=True)
+            dx = dx_gn
+        elif t * numpy.linalg.norm(dx_sd) > max(1.0, numpy.linalg.norm(xold)) * (c ** microiter): # GN step outside, SD step also outside
+            print('  Trust Region Optimization Step ', microiter, ': Steepest Descent', flush=True)
+            dx = (c ** microiter) / numpy.linalg.norm(dx_sd) * dx_sd
+        else: # GN step outside, SD step inside (dog leg step)
+              # dx := t*dx_sd + s (dx_gn - t*dx_sd) s.t. ||dx|| = c^p
+            print('  Trust Region Optimization Step ', microiter, ': Dog Leg', flush=True)
+            tdx_sd = t*dx_sd
+            diff = dx_gn - tdx_sd
+            #s = (-dx_sd.T@diff + numpy.sqrt((dx_sd.T@diff)**2 - numpy.linalg.norm(diff)**2*(numpy.linalg.norm(dx_sd)**2-(c ** microiter)**2))) / (numpy.linalg.norm(dx_sd))**2
+            # s is largest value in [0, 1] s.t. ||dx|| \le trust radius
+            s = 1
+            dx = tdx_sd + s*diff
+            while (numpy.linalg.norm(dx) > c ** microiter and s > 0):
+                s -= 0.001
+                dx = tdx_sd + s*diff
+        if prevdx is None or not numpy.all(dx == prevdx):
+            # Actual Reduction := f(x_k) - f(x_k + dx)
+            fnew = func(xold + dx)
+            ared = 0.5 * (numpy.linalg.norm(fold)**2 - numpy.linalg.norm(fnew)**2)
+            # Predicted Reduction := q(0) - q(dx) where q = (1/2) ||F_k + B_k d||^2
+            pred = 0.5 * (numpy.linalg.norm(fold)**2 - numpy.linalg.norm(fold + B@dx)**2)
+        # Trust Region convergence criteria
+        # r = ared/pred \le rho
+        ratio = ared / pred
+        microiter += 1
+        if prevdx is None or not numpy.all(dx == prevdx) and be_var.PRINT_LEVEL > 2:
+            print('    ||δx||: ', numpy.linalg.norm(dx), flush=True)
+            print('    Reduction Ratio (Actual / Predicted): ', ared, '/', pred, '=', ratio, flush=True)
+        prevdx = dx
+    return xold + dx, fnew # xnew
 
 class FrankQN:
     """ Quasi Newton Optimization
@@ -59,7 +127,7 @@ class FrankQN:
 
     """
 
-    def __init__(self, func, x0, f0, J0, trust=0.3, max_space=500):
+    def __init__(self, func, x0, f0, J0, trust=0.5, max_space=500):
 
         self.x0 = x0
         self.n = x0.size        
@@ -83,7 +151,7 @@ def __init__(self, func, x0, f0, J0, trust=0.3, max_space=500):
         self.B = None
         self.trust = trust
 
-    def next_step(self):
+    def next_step(self, trust_region=False):
 
         if self.iter_ == 0:            
             self.xnew = self.x0            
@@ -106,14 +174,17 @@ def next_step(self):
                 (dx_i@self.Binv@df_i)
             self.Binv += tmp__            
             us_tmp = self.Binv@self.fnew
-
-        self.us[self.iter_] = self.get_Bnfn(self.iter_)
-
-        alp, self.xnew, self.fnew = line_search_LF(
-            self.func, self.xold, self.fold, -self.us[self.iter_], self.iter_)
-
-        # udpate vs, dxs, and fs        
-        self.vs[self.iter_] = numpy.dot(self.B0, self.fnew)
+
+        if trust_region:
+            self.xnew, self.fnew = trustRegion(self.func, self.xold, self.fold, self.Binv, c = self.trust)
+        else:
+            self.us[self.iter_] = self.get_Bnfn(self.iter_)
+
+            alp, self.xnew, self.fnew = line_search_LF(
+                self.func, self.xold, self.fold, -self.us[self.iter_], self.iter_)
+
+            # udpate vs, dxs, and fs        
+            self.vs[self.iter_] = numpy.dot(self.B0, self.fnew)
         self.dxs[self.iter_] = self.xnew - self.xold
         self.fs[self.iter_+1] = self.fnew.copy()
 

tests/dm_molBE_test.py

@@ -0,0 +1,34 @@
+"""
+This script tests the correlation energies of sample restricted molecular BE calculations from density matching
+
+Author(s): Minsik Cho
+"""
+
+import os, unittest
+from pyscf import cc, gto, scf
+from molbe import fragpart, BE
+
+class TestBE_restricted(unittest.TestCase):
+    # TODO: Add test against known values (molecular_restrict_test)
+    def test_h8_sto3g_ben_trustRegion(self):
+        # Test consistency between two QN methods
+        mol = gto.M()
+        mol.atom = [['H', (0.,0.,i)] for i in range(7)]; mol.atom.append(['H', (0.,0.,4.2)])
+        mol.basis = 'sto-3g'
+        mol.charge = 0.; mol.spin = 0.
+        mol.build()
+        self.molecular_QN_test(mol, 'be2', 'H8 (BE2)', 'hchain_simple', only_chem = False)
+
+    def molecular_QN_test(self, mol, be_type, test_name, frag_type, delta = 1e-6, only_chem = True):
+        mf = scf.RHF(mol); mf.max_cycle = 100; mf.kernel()
+        fobj = fragpart(frag_type=frag_type, be_type=be_type, mol = mol)
+        mybe1 = BE(mf, fobj)
+        mybe1.optimize(solver='CCSD', method='QN', trust_region=False, only_chem=only_chem)
+        mybe2 = BE(mf, fobj)
+        mybe2.optimize(solver='CCSD', method='QN', trust_region=True, only_chem=only_chem)
+        self.assertAlmostEqual(mybe1.ebe_tot, mybe2.ebe_tot,
+                               msg = "BE Correlation Energy (DM) for " + test_name
+                               + " does not return comparable results from two QN optimization!", delta = delta)
+
+if __name__ == '__main__':
+    unittest.main()