UBE update

example/molbe_oneshot_ube_qmmm.py

@@ -33,7 +33,8 @@
     unrestricted=True,  # specify unrestricted calculation
     from_chk=False,  # can save the UHF as PySCF checkpoint.
     # Set to true if running from converged UHF chk
-    checkfile=None,
-)  # if not None, will save UHF calculation to a checkfile.
+    checkfile=None,  # if not None, will save UHF calculation to checkfile
+    opt="SOSCF",  # SOSCF and DAMP options: note that damping settings hard-coded
+)
 # if rerunning from chk (from_chk=True), name the checkfile here
-#            ecp = ecp) # can add ECP for heavy atoms as: {'Ru': 'def2-SVP'}
+#  ecp = ecp) # can add ECP for heavy atoms as: {'Ru': 'def2-SVP'}

pyproject.toml

@@ -46,3 +46,6 @@ where = ["./src"]
 include = ["quemb*"]
 exclude = []
 namespaces = false
+
+[project.scripts]
+ube-hfcc = "quemb.scripts.ube_hfcc:main"

src/quemb/molbe/be_parallel.py

@@ -407,7 +407,7 @@ def run_solver_u(
         gcores=full_uhf.full_gcore,
         frozen=frozen,
     )
-    return e_f
+    return e_f, rdm1_tmp, fobj_a._mf.mo_coeff.copy(), fobj_b._mf.mo_coeff.copy()
 
 
 def be_func_parallel(
@@ -473,7 +473,7 @@ def be_func_parallel(
         the error norm, error vector, and the computed energy.
     """
     # Set the number of OpenMP threads
-    os.system("export OMP_NUM_THREADS=" + str(ompnum))
+    os.environ["OMP_NUM_THREADS"] = str(ompnum)
     nprocs = nproc // ompnum
 
     # Update the effective Hamiltonian with potentials
@@ -603,7 +603,7 @@ def be_func_parallel_u(
         Returns the computed energy
     """
     # Set the number of OpenMP threads
-    os.system("export OMP_NUM_THREADS=" + str(ompnum))
+    os.environ["OMP_NUM_THREADS"] = str(ompnum)
     nprocs = nproc // ompnum
 
     with ProcessPool(nprocs) as pool_:
@@ -623,14 +623,21 @@ def be_func_parallel_u(
             )
             results.append(result)
 
-        energy_list = [result.get() for result in results]
-
+        results_list = [result.get() for result in results]
+    # Store RDMs back into fragment objects
+    for i, (fobj_a, fobj_b) in enumerate(Fobjs):
+        e_f, rdm1_tmp, mo_a, mo_b = results_list[i]
+        fobj_a.rdm1__ = rdm1_tmp[0].copy()
+        fobj_b.rdm1__ = rdm1_tmp[1].copy()
+        fobj_a.mo_coeffs = mo_a
+        fobj_b.mo_coeffs = mo_b
     # Compute and return fragment energy
     e_1 = 0.0
     e_2 = 0.0
     e_c = 0.0
-    for i in range(len(energy_list)):
-        e_1 += energy_list[i][0]
-        e_2 += energy_list[i][1]
-        e_c += energy_list[i][2]
+    for i in range(len(results_list)):
+        e_f = results_list[i][0]
+        e_1 += e_f[0]
+        e_2 += e_f[1]
+        e_c += e_f[2]
     return (e_1 + e_2 + e_c, (e_1, e_2, e_c))

src/quemb/molbe/mbe.py

@@ -1422,7 +1422,7 @@ def localize(
                     Cpop = multi_dot((C_.T, self.S, C_))
                     no_core_idx = where(diag(Cpop) > 0.7)[0]
                     C_ = C_[:, no_core_idx]
-                    S_ = multi_dot((C_.T, self.S, C_))
+                    S_ = multi_dot((C_.T, self.S, C_))  # type: ignore[assignment]
                     es_, vs_ = eigh(S_)
                     s_ = sqrt(es_)
                     s_ = diag(1.0 / s_)
@@ -1459,7 +1459,7 @@ def localize(
                 Cpop = diag(Cpop)
                 no_core_idx = where(Cpop > 0.55)[0]
                 C_ = C_[:, no_core_idx]
-                S_ = multi_dot((C_.T, self.S, C_))
+                S_ = multi_dot((C_.T, self.S, C_))  # type: ignore[assignment]
                 es_, vs_ = eigh(S_)
                 s_ = sqrt(es_)
                 s_ = diag(1.0 / s_)

src/quemb/molbe/misc.py

@@ -265,6 +265,7 @@ def be2puffin(
     checkfile=None,
     ecp=None,
     frag_type="chemgen",
+    opt="soscf",
 ):
     """Front-facing API bridge tailored for SCINE Puffin
 
@@ -369,36 +370,56 @@ def be2puffin(
                         "Using QM/MM Point Charges: Assuming QM structure in Angstrom "
                         "and MM Coordinates in Bohr !!!"
                     )
-                    mf1 = scf.UHF(mol).set(
-                        max_cycle=200
-                    )  # using SOSCF is more reliable
-                    # mf1 = scf.UHF(mol).set(max_cycle = 200, level_shift = (0.3, 0.2))
-                    # using level shift helps, but not always. level_shift and
-                    # scf.addons.dynamic_level_shift do not seem to work with QM/MM
-                    # note: from the SCINE database, the structure is in Angstrom but
-                    # the MM point charges are in Bohr !!
-                    mf = qmmm.mm_charge(
-                        mf1, pts_and_charges[0], pts_and_charges[1], unit="bohr"
-                    ).newton()  # mf object, coordinates, charges
+                    mf1 = scf.UHF(mol).set(max_cycle=200)
+                    if opt.upper() == "SOSCF":
+                        # using SOSCF is more reliable
+                        # note: from the SCINE database, the structure is in Angstrom
+                        # but the MM point charges are in Bohr !!
+                        mf = qmmm.mm_charge(
+                            mf1, pts_and_charges[0], pts_and_charges[1], unit="bohr"
+                        ).newton()  # mf object, coordinates, charges
+                    elif opt.upper() == "DAMP":
+                        mf = qmmm.mm_charge(
+                            mf1, pts_and_charges[0], pts_and_charges[1], unit="bohr"
+                        )
+                        mf.damp = 0.5
+                        mf.diis_start_cycle = 5
                 else:
-                    mf = scf.UHF(mol).set(max_cycle=200, level_shift=(0.3, 0.2))
+                    if opt.upper() == "SOSCF":
+                        mf = scf.UHF(mol).set(max_cycle=200).newton()
+                    elif opt.upper() == "DAMP":
+                        mf = scf.UHF(mol).set(max_cycle=200)
+                        mf.damp = 0.5
+                        mf.diis_start_cycle = 5
             else:
-                mf = scf.UHF(mol).set(max_cycle=200).newton()
+                if opt.upper() == "SOSCF":
+                    mf = scf.UHF(mol).set(max_cycle=200).newton()
+                elif opt.upper() == "DAMP":
+                    mf = scf.UHF(mol).set(max_cycle=200)
+                    mf.damp = 0.5
+                    mf.diis_start_cycle = 5
         else:  # restricted
             if pts_and_charges:  # running QM/MM
                 print(
                     "Using QM/MM Point Charges: Assuming QM structure in Angstrom and "
                     "MM Coordinates in Bohr !!!"
                 )
-                mf1 = scf.RHF(mol).set(max_cycle=200)
-                mf = qmmm.mm_charge(
-                    mf1, pts_and_charges[0], pts_and_charges[1], unit="bohr"
-                ).newton()
                 if use_df or jk is not None:
                     raise ValueError(
                         "Setting use_df to false and jk to none: have not tested DF "
                         "and QM/MM from point charges at the same time"
                     )
+                mf1 = scf.RHF(mol).set(max_cycle=200)
+                if opt.upper() == "SOSCF":
+                    mf = qmmm.mm_charge(
+                        mf1, pts_and_charges[0], pts_and_charges[1], unit="bohr"
+                    ).newton()
+                elif opt.upper() == "DAMP":
+                    mf = qmmm.mm_charge(
+                        mf1, pts_and_charges[0], pts_and_charges[1], unit="bohr"
+                    )
+                    mf.damp = 0.5
+                    mf.diis_start_cycle = 5
             elif use_df and jk is None:
                 mf = scf.RHF(mol).density_fit(auxbasis=df_aux_basis)
             else:
@@ -464,12 +485,22 @@ def be2puffin(
                 "Using QM/MM Point Charges: Assuming QM structure in Angstrom and "
                 "MM Coordinates in Bohr !!!"
             )
-            mf = qmmm.mm_charge(
-                mf,
-                pts_and_charges[0],
-                pts_and_charges[1],
-                unit="bohr",
-            ).newton()
+            if opt.upper() == "SOSCF":
+                mf = qmmm.mm_charge(
+                    mf,
+                    pts_and_charges[0],
+                    pts_and_charges[1],
+                    unit="bohr",
+                ).newton()
+            elif opt.upper() == "DAMP":
+                mf = qmmm.mm_charge(
+                    mf,
+                    pts_and_charges[0],
+                    pts_and_charges[1],
+                    unit="bohr",
+                )
+                mf.damp = 0.5
+                mf.diis_start_cycle = 5
         time_post_mf = time.time()
         print("Chkfile electronic energy:", mf.energy_elec(), flush=True)
         print("Chkfile e_tot:", mf.e_tot, flush=True)
@@ -496,7 +527,10 @@ def be2puffin(
     # Run oneshot embedding and return system energy
 
     mybe.oneshot(solver=solver, nproc=nproc, ompnum=ompnum)
-    return mybe.ebe_tot - mybe.ebe_hf
+    if unrestricted:
+        return mybe.ebe_tot - mybe.hf_etot
+    else:
+        return mybe.ebe_tot - mybe.ebe_hf
 
 
 def print_energy_cumulant(ecorr, e_V_Kapprox, e_F_dg, e_hf):

src/quemb/molbe/pfrag.py

@@ -6,7 +6,6 @@
 import numpy as np
 import scipy.linalg
 from numpy import (
-    argsort,
     array,
     diag_indices,
     einsum,
@@ -469,28 +468,36 @@ def schmidt_decomposition(
 
     # Identify significant environment orbitals based on eigenvalue threshold
     Bidx = []
-
+    for i in range(len(Eval)):
+        if thr_bath < np.abs(Eval[i]) < 1.0 - thr_bath:
+            Bidx.append(i)
     # Set the number of orbitals to be taken from the environment orbitals
     # Based on an eigenvalue threshold ordering
     if norb is not None:
-        n_frag_ind = len(Frag_sites1)
-        n_bath_ind = norb - n_frag_ind
-        ind_sort = argsort(np.abs(Eval))
-        first_el = [x for x in ind_sort if x < 1.0 - thr_bath][-1 * n_bath_ind]
-        for i in range(len(Eval)):
-            if np.abs(Eval[i]) >= first_el:
-                Bidx.append(i)
-    else:
-        for i in range(len(Eval)):
-            if thr_bath < np.abs(Eval[i]) < 1.0 - thr_bath:
-                Bidx.append(i)
+        # add extra orbital(s) from the environment; these will likely have
+        # Eval close to 1. note: there are normally very few orbitals with a
+        # Eval[i] <= thr_bath, so adding Bidx from the "front of the list"
+        # doesn't work. Instead, we add the excluded orbitals closest to the
+        # thr_bath/1-thr_bath boundary (this is analagous to tightening up
+        # the bath threshold for the alpha or beta orbitals until they are
+        # the same size)
+        excluded = [i for i in range(len(Eval)) if i not in set(Bidx)]
+        excluded_sorted = sorted(
+            excluded,
+            key=lambda i: min(abs(Eval[i] - (1.0 - thr_bath)), abs(Eval[i] - thr_bath)),
+        )
+        # Bidx corresponds to sorted Eval and Evec, so this adds indices
+        # closest to the bath threshold until the bath size reaches norb
+        for idx in excluded_sorted:
+            if len(Bidx) >= norb:
+                break
+            Bidx.append(idx)
 
     # Initialize the transformation matrix (TA)
     TA = zeros([Tot_sites, len(AO_in_frag) + len(Bidx)])
     TA[AO_in_frag, : len(AO_in_frag)] = eye(len(AO_in_frag))  # Fragment part
     TA[Env_sites1, len(AO_in_frag) :] = Evec[:, Bidx]  # Environment part
 
-    # return TA, norbs_frag, norbs_bath
     return TA, Frag_sites1.shape[0], len(Bidx)
 
 

src/quemb/molbe/solver.py

@@ -596,10 +596,6 @@ def be_func_u(
         Whether to evaluate the energy. Defaults to False.
     relax_density : bool, optional
         Whether to relax the density. Defaults to False.
-    return_vec : bool, optional
-        Whether to return the error vector. Defaults to False.
-    ebe_hf : float, optional
-        Hartree-Fock energy. Defaults to 0.
     use_cumulant : bool, optional
         Whether to use the cumulant-based energy expression. Defaults to True.
     frozen : bool, optional
@@ -635,10 +631,9 @@ def be_func_u(
 
         assert fobj_a._mf is not None and fobj_b._mf is not None
         fobj_a.rdm1__ = rdm1_tmp[0].copy()
-        fobj_b._rdm1 = (
+        fobj_a._rdm1 = (
             multi_dot((fobj_a._mf.mo_coeff, rdm1_tmp[0], fobj_a._mf.mo_coeff.T)) * 0.5
         )
-
         fobj_b.rdm1__ = rdm1_tmp[1].copy()
         fobj_b._rdm1 = (
             multi_dot((fobj_b._mf.mo_coeff, rdm1_tmp[1], fobj_b._mf.mo_coeff.T)) * 0.5