DMRG Fragment Solver and Associated Utilities (#41)

* Expose pipek-mezey lo modules to `lo_method`

* DMRG fragment solver via `block2`.

* Change default scratch directory for `solver='SHCI'`

* Expose DMRG to `oneshot`, cleanup scratch m'mt

* Added BE-DMRG example and `examples/figures/`

* Update `scratch` keywords to use `scratch_dir` by default

---------

Co-authored-by: Shaun Weatherly <shaunrw@telemachus.mit.edu>

example/molbe_dmrg_block2.py

@@ -0,0 +1,142 @@
+# A run through the `QuEmb` interface with `block2` for performing BE-DMRG.
+# `block2` is a DMRG and sparse tensor network library developed by the 
+# Garnet-Chan group at Caltech: https://block2.readthedocs.io/en/latest/index.html
+
+import os, numpy, sys
+import matplotlib.pyplot as plt
+from pyscf import gto, scf, fci, cc
+from molbe import BE, fragpart
+
+# We'll consider the dissociation curve for a 1D chain of 8 H-atoms:
+num_points = 3
+seps = numpy.linspace(0.60, 1.6, num=num_points)
+fci_ecorr, ccsd_ecorr, ccsdt_ecorr, bedmrg_ecorr = [], [], [], []
+
+# Specify a scratch directory for fragment DMRG files:
+scratch = os.getcwd() 
+
+for a in seps:
+    # Hartree-Fock serves as the starting point for all BE calculations:
+    mol = gto.M()
+    mol.atom = [['H', (0.,0.,i*a)] for i in range(8)]
+    mol.basis = 'sto-3g'
+    mol.charge = 0
+    mol.spin = 0
+    mol.build()
+
+    mf = scf.RHF(mol)
+    mf.conv_tol = 1e-12
+    mf.kernel()
+
+    # Exact diagonalization (FCI) will provide the reference:
+    mc = fci.FCI(mf)
+    fci_ecorr.append(mc.kernel()[0] - mf.e_tot)
+
+    # CCSD and CCSD(T) are good additional points of comparison:
+    mycc = cc.CCSD(mf).run()
+    et_correction = mycc.ccsd_t()
+    ccsd_ecorr.append(mycc.e_tot - mf.e_tot)
+    ccsdt_ecorr.append(mycc.e_tot + et_correction - mf.e_tot)
+
+    # Define BE1 fragments. Prior to DMRG the localization of MOs 
+    # is usually necessary. While there doesn't appear to be 
+    # any clear advantage to using any one scheme over another,
+    # the Pipek-Mezey scheme continues to be the most popular. With
+    # BE-DMRG, localization takes place prior to fragmentation:
+    fobj = fragpart(be_type='be1', mol=mol)
+    mybe = BE(
+        mf, 
+        fobj, 
+        lo_method='pipek-mezey',                        # or 'lowdin', 'iao', 'boys'
+        pop_method='lowdin'                             # or 'meta-lowdin', 'mulliken', 'iao', 'becke'
+        )
+
+    # Next, run BE-DMRG with default parameters and maxM=100.
+    mybe.oneshot(
+        solver='block2',                                # or 'DMRG', 'DMRGSCF', 'DMRGCI'
+        scratch=scratch,                                # Scratch dir for fragment DMRG
+        maxM=100,                                       # Max fragment bond dimension
+        force_cleanup=True,                             # Remove all fragment DMRG tmpfiles
+        ) 
+
+    bedmrg_ecorr.append(mybe.ebe_tot - mf.e_tot)
+    # Setting `force_cleanup=True` will clean the scratch directory after each 
+    # fragment DMRG calculation finishes. DMRG tempfiles can be quite large, so
+    # be sure to keep an eye on them if `force_cleanup=False` (default).
+
+    # NOTE: This does *not* delete the log files `dmrg.conf` and `dmrg.out`for each frag,
+    # which can still be found in `/scratch/`.
+
+# Finally, plot the resulting potential energy curves:
+fig, ax = plt.subplots()
+
+ax.plot(seps, fci_ecorr, 'o-', linewidth=1, label='FCI')
+ax.plot(seps, ccsd_ecorr, 'o-', linewidth=1, label='CCSD')
+ax.plot(seps, ccsdt_ecorr, 'o-', linewidth=1, label='CCSD(T)')
+ax.plot(seps, bedmrg_ecorr, 'o-', linewidth=1, label='BE1-DMRG')
+ax.legend()
+
+plt.savefig(os.path.join(scratch, f'BEDMRG_H8_PES{num_points}.png'))
+
+# (See ../quemb/example/figures/BEDMRG_H8_PES20.png for an example.)
+
+# For larger fragments, you'll want greater control over the fragment
+# DMRG calculations. Using the same setup as above for a single geometry: 
+mol = gto.M()
+mol.atom = [['H', (0.,0.,i*1.2)] for i in range(8)]
+mol.basis = 'sto-3g'
+mol.charge = 0
+mol.spin = 0
+mol.build()
+fobj = fragpart(be_type='be2', mol=mol)
+mybe = BE(
+    mf, 
+    fobj, 
+    lo_method='pipek-mezey',
+    pop_method='lowdin'
+    )
+
+# We automatically construct the fragment DMRG schedules based on user keywords. The following
+# input, for example, yields a 60 sweep schedule which uses the two-dot algorithm from sweeps 0-49,
+# and the one-dot algo from 50-60. The fragment orbitals are also reordered according the Fiedler 
+# vector procedure, along with a few other tweaks:
+
+mybe.optimize(
+    solver='block2',                                # or 'DMRG', 'DMRGSCF', 'DMRGCI'
+    scratch=scratch,                                # Scratch dir for fragment DMRG
+    startM=20,                                      # Initial fragment bond dimension (1st sweep)
+    maxM=200,                                       # Maximum fragment bond dimension
+    max_iter=60,                                    # Max number of sweeps
+    twodot_to_onedot=50,                            # Sweep num to switch from two- to one-dot algo.
+    max_mem=40,                                     # Max memory (in GB) allotted to fragment DMRG
+    max_noise=1e-3,                                 # Max MPS noise introduced per sweep
+    min_tol=1e-8,                                   # Tighest Davidson tolerance per sweep
+    block_extra_keyword=['fiedler'],                # Specify orbital reordering algorithm
+    force_cleanup=True,                             # Remove all fragment DMRG tmpfiles
+    only_chem=True,
+) 
+
+# Or, alternatively, we can construct a full schedule by hand:
+schedule={
+    'scheduleSweeps': [0, 10, 20, 30, 40, 50],                  # Sweep indices
+    'scheduleMaxMs': [25, 50, 100, 200, 500, 500],              # Sweep maxMs
+    'scheduleTols': [1e-5,1e-5, 1e-6, 1e-6, 1e-8, 1e-8],        # Sweep Davidson tolerances
+    'scheduleNoises': [0.01, 0.01, 0.001, 0.001, 1e-4, 0.0],    # Sweep MPS noise
+}
+
+# and pass it to the fragment solver through `schedule_kwargs`:
+mybe.optimize(
+    solver='block2',
+    scratch=scratch, 
+    schedule_kwargs=schedule,
+    block_extra_keyword=['fiedler'],
+    force_cleanup=True,
+    only_chem=True,
+)
+
+# To make sure the calculation is proceeding as expected, make sure to check `[scratch]/dmrg.conf` 
+# and `[scratch]/dmrg.out`, which are the fragment DMRG inputs and outputs, respectively, used
+# by `block2`. 
+
+#NOTE: Parameters in `schedule_kwargs` will overwrite any other DMRG kwargs.
+#NOTE: The DMRG schedule kwargs and related syntax follows the standard notation used in block2.

molbe/_opt.py

@@ -45,7 +45,7 @@ def __init__(self, pot, Fobjs, Nocc, enuc,solver='MP2', ecore=0.,
                  only_chem=False, hf_veff = None,
                  hci_pt=False,hci_cutoff=0.001, ci_coeff_cutoff = None, select_cutoff=None,
                  max_space=500, conv_tol = 1.e-6,relax_density = False,
-                 ebe_hf =0.):
+                 ebe_hf =0., scratch_dir=None, **solver_kwargs):
 
         # Initialize class attributes 
         self.ebe_hf=ebe_hf
@@ -70,6 +70,8 @@ def __init__(self, pot, Fobjs, Nocc, enuc,solver='MP2', ecore=0.,
         self.ci_coeff_cutoff = ci_coeff_cutoff
         self.select_cutoff = select_cutoff
         self.hci_pt = hci_pt
+        self.solver_kwargs=solver_kwargs
+        self.scratch_dir=scratch_dir
 
     def objfunc(self, xk):
         """
@@ -97,7 +99,8 @@ def objfunc(self, xk):
                                          nproc=self.ompnum, relax_density=self.relax_density,
                                          ci_coeff_cutoff = self.ci_coeff_cutoff,
                                          select_cutoff = self.select_cutoff, hci_pt=self.hci_pt,
-                                         ecore=self.ecore, ebe_hf=self.ebe_hf, be_iter=self.iter)
+                                         ecore=self.ecore, ebe_hf=self.ebe_hf, be_iter=self.iter,
+                                         scratch_dir=self.scratch_dir, **self.solver_kwargs)
         else:
             err_, errvec_,ebe_ = be_func_parallel(xk, self.Fobjs, self.Nocc, self.solver, self.enuc,
                                                   eeval=True, return_vec=True, hf_veff = self.hf_veff,
@@ -106,7 +109,8 @@ def objfunc(self, xk):
                                                   hci_cutoff=self.hci_cutoff,relax_density=self.relax_density,
                                                   ci_coeff_cutoff = self.ci_coeff_cutoff,
                                                   select_cutoff = self.select_cutoff,
-                                                  ecore=self.ecore, ebe_hf=self.ebe_hf, be_iter=self.iter)
+                                                  ecore=self.ecore, ebe_hf=self.ebe_hf, be_iter=self.iter, 
+                                                  scratch_dir=self.scratch_dir, **self.solver_kwargs)
 
         # Update error and BE energy
         self.err = err_
@@ -173,7 +177,7 @@ 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):
+             J0=None, nproc=1, ompnum=4, max_iter=500, scratch_dir=None,  **solver_kwargs):
     """BE optimization function
 
     Interfaces BEOPT to perform bootstrap embedding optimization.
@@ -216,13 +220,14 @@ def optimize(self, solver='MP2',method='QN',
 
     # Initialize the BEOPT object
     be_ = BEOPT(pot, self.Fobjs, self.Nocc, self.enuc, hf_veff = self.hf_veff,
-                nproc=nproc, ompnum=ompnum,
+                nproc=nproc, ompnum=ompnum, scratch_dir=scratch_dir,
                 max_space=max_iter,conv_tol = conv_tol,
                 only_chem=only_chem,
                 hci_cutoff=self.hci_cutoff,
                 ci_coeff_cutoff = self.ci_coeff_cutoff,relax_density=relax_density,
                 select_cutoff = self.select_cutoff,hci_pt=self.hci_pt,
-                solver=solver, ecore=self.E_core, ebe_hf=self.ebe_hf)
+                solver=solver, ecore=self.E_core, ebe_hf=self.ebe_hf,
+                **solver_kwargs)
 
     if method=='QN':
         # Prepare the initial Jacobian matrix

molbe/fragment.py

@@ -108,6 +108,7 @@ def hchain_simple(self):
         if self.be_type=='be1':
             for i in range(self.natom):
                 self.fsites.append([i])
+                self.edge.append([])
             self.Nfrag = len(self.fsites)
 
         elif self.be_type=='be2':

molbe/lo.py

@@ -1,5 +1,4 @@
-# Author(s): Henry Tran
-#            Oinam Meitei
+# Author(s): Henry Tran, Oinam Meitei, Shaun Weatherly
 #
 from pyscf import lib
 import numpy,sys
@@ -193,7 +192,7 @@ def get_pao_native(Ciao, S, mol, valence_basis):
 
     return Cpao
 
-def get_loc(mol, C, method):
+def get_loc(mol, C, method, pop_method=None, init_guess=None):
     if method.upper() == 'ER':
         from pyscf.lo import ER as Localizer
     elif method.upper() == 'PM':
@@ -202,19 +201,18 @@ def get_loc(mol, C, method):
         from pyscf.lo import Boys as Localizer
     else:
         raise NotImplementedError('Localization scheme not understood')
-
+    
     mlo = Localizer(mol, C)
-    mlo.init_guess = None 
+    if pop_method is not None:
+        mlo.pop_method=str(pop_method)
+
+    mlo.init_guess = init_guess 
     C_ = mlo.kernel()
 
     return C_
 
-
-
-
-
 def localize(self, lo_method, mol=None, valence_basis='sto-3g',
-             hstack=False,
+             hstack=False, pop_method=None, init_guess=None,
              valence_only=False, nosave=False):
     """ Molecular orbital localization
 
@@ -298,6 +296,37 @@ def localize(self, lo_method, mol=None, valence_basis='sto-3g',
                 self.lmo_coeff = functools.reduce(numpy.dot,
                                                   (self.W.T, self.S, self.C))            
 
+    elif lo_method in ['pipek-mezey','pipek', 'PM']:
+
+        es_, vs_ = eigh(self.S)
+        edx = es_ > 1.e-15                
+        self.W = numpy.dot(vs_[:,edx]/numpy.sqrt(es_[edx]), vs_[:,edx].T)
+
+        es_, vs_ = eigh(self.S)
+        edx = es_ > 1.e-15                
+        W_ = numpy.dot(vs_[:,edx]/numpy.sqrt(es_[edx]), vs_[:,edx].T)
+        if self.frozen_core:                    
+            P_core = numpy.eye(W_.shape[0]) - numpy.dot(self.P_core, self.S)
+            C_ = numpy.dot(P_core, W_)
+            Cpop = functools.reduce(numpy.dot,
+                                    (C_.T, self.S, C_))
+            Cpop = numpy.diag(Cpop)
+            no_core_idx = numpy.where(Cpop > 0.55)[0]
+            C_ = C_[:,no_core_idx]
+            S_ = functools.reduce(numpy.dot, (C_.T, self.S, C_))
+            es_, vs_ = eigh(S_)
+            s_ = numpy.sqrt(es_)
+            s_ = numpy.diag(1.0/s_)
+            W_ = functools.reduce(numpy.dot,
+                                    (vs_, s_, vs_.T))
+            W_ = numpy.dot(C_, W_)            
+
+        self.W = get_loc(self.mol, W_, 'PM', pop_method=pop_method, init_guess=init_guess)
+
+        if not self.frozen_core:
+            self.lmo_coeff = self.W.T @ self.S @ self.C
+        else:                
+            self.lmo_coeff = self.W.T @ self.S @ self.C[:,self.ncore:]
 
     elif lo_method=='iao':        
         from pyscf import lo

molbe/mbe.py

@@ -49,12 +49,12 @@ class BE:
     """
 
     def __init__(self, mf, fobj, eri_file='eri_file.h5', 
-                 lo_method='lowdin',compute_hf=True, 
+                 lo_method='lowdin', pop_method=None, compute_hf=True, 
                  restart=False, save=False,
                  restart_file='storebe.pk',
                  mo_energy = None, 
                  save_file='storebe.pk',hci_pt=False,
-                 nproc=1, ompnum=4,
+                 nproc=1, ompnum=4, scratch_dir=None,
                  hci_cutoff=0.001, ci_coeff_cutoff = None, select_cutoff=None,
                  integral_direct_DF=False, auxbasis = None):
         """
@@ -163,6 +163,7 @@ def __init__(self, mf, fobj, eri_file='eri_file.h5',
         self.Fobjs = []
         self.pot = initialize_pot(self.Nfrag, self.edge_idx)
         self.eri_file = eri_file
+        self.scratch_dir = scratch_dir
 
         # Set scratch directory
         jobid=''
@@ -173,7 +174,7 @@ def __init__(self, mf, fobj, eri_file='eri_file.h5',
                 jobid = ''
         if not be_var.SCRATCH=='': 
             self.scratch_dir = be_var.SCRATCH+str(jobid)
-            os.system('mkdir '+self.scratch_dir)
+            os.system('mkdir -p '+self.scratch_dir)
         else:
             self.scratch_dir = None
         if jobid == '':
@@ -208,10 +209,10 @@ def __init__(self, mf, fobj, eri_file='eri_file.h5',
 
         if not restart:
             # Localize orbitals
-            self.localize(lo_method, mol=self.mol, valence_basis=fobj.valence_basis, valence_only=fobj.valence_only)
+            self.localize(lo_method, pop_method=pop_method, mol=self.mol, valence_basis=fobj.valence_basis, valence_only=fobj.valence_only)
 
             if fobj.valence_only and lo_method=='iao':
-                self.Ciao_pao = self.localize(lo_method, mol=self.mol, valence_basis=fobj.valence_basis,
+                self.Ciao_pao = self.localize(lo_method, pop_method=pop_method, mol=self.mol, valence_basis=fobj.valence_basis,
                                               hstack=True,
                                               valence_only=False, nosave=True)
 
@@ -375,7 +376,8 @@ def initialize(self, eri_,compute_hf, restart=False):
             fobj.udim = couti
             couti = fobj.set_udim(couti)
 
-    def oneshot(self, solver='MP2', nproc=1, ompnum=4, calc_frag_energy=False, clean_eri=False):
+    def oneshot(self, solver='MP2', nproc=1, ompnum=4, calc_frag_energy=False, clean_eri=False, 
+                scratch_dir=None, **solver_kwargs):
         """
         Perform a one-shot bootstrap embedding calculation.
 
@@ -394,22 +396,25 @@ def oneshot(self, solver='MP2', nproc=1, ompnum=4, calc_frag_energy=False, clean
         """
         from .solver import be_func
         from .be_parallel import be_func_parallel
-
+
+        self.scratch_dir = scratch_dir
+        self.solver_kwargs = solver_kwargs
+
         print("Calculating Energy by Fragment? ", calc_frag_energy)
         if nproc == 1:
             rets  = be_func(None, self.Fobjs, self.Nocc, solver, self.enuc, hf_veff=self.hf_veff,
                         hci_cutoff=self.hci_cutoff,
                         ci_coeff_cutoff = self.ci_coeff_cutoff,
                         select_cutoff = self.select_cutoff,
                         nproc=ompnum, frag_energy=calc_frag_energy,
-                        ereturn=True, eeval=True)
+                        ereturn=True, eeval=True, scratch_dir=self.scratch_dir, **self.solver_kwargs)
         else:
             rets  = be_func_parallel(None, self.Fobjs, self.Nocc, solver, self.enuc, hf_veff=self.hf_veff,
                                  hci_cutoff=self.hci_cutoff,
                                  ci_coeff_cutoff = self.ci_coeff_cutoff,
                                  select_cutoff = self.select_cutoff,
                                  ereturn=True, eeval=True, frag_energy=calc_frag_energy,
-                                 nproc=nproc, ompnum=ompnum)
+                                 nproc=nproc, ompnum=ompnum, scratch_dir=self.scratch_dir, **self.solver_kwargs)
 
         print('-----------------------------------------------------',
                   flush=True)