MOM for GHF and DHF (fix issue pyscf#1461)

sichaox · Mar 23, 2023 · 98f232f · 98f232f
1 parent ec8618c
commit 98f232f
Showing 1 changed file with 52 additions and 37 deletions.
diff --git a/pyscf/scf/addons.py b/pyscf/scf/addons.py
@@ -268,9 +268,9 @@ def get_occ(mo_energy, mo_coeff=None):
 
 def mom_occ_(mf, occorb, setocc):
     '''Use maximum overlap method to determine occupation number for each orbital in every
-    iteration. It can be applied to unrestricted HF/KS and restricted open-shell
-    HF/KS.'''
+    iteration.'''
     from pyscf.scf import uhf, rohf
+    log = logger.Logger(mf.stdout, mf.verbose)
     if isinstance(mf, uhf.UHF):
         coef_occ_a = occorb[0][:, setocc[0]>0]
         coef_occ_b = occorb[1][:, setocc[1]>0]
@@ -279,44 +279,59 @@ def mom_occ_(mf, occorb, setocc):
             raise ValueError('Wrong occupation setting for restricted open-shell calculation.')
         coef_occ_a = occorb[:, setocc[0]>0]
         coef_occ_b = occorb[:, setocc[1]>0]
+    else: # GHF, and DHF
+        assert setocc.ndim == 1
+
+    if isinstance(mf, (uhf.UHF, rohf.ROHF)):
+        def get_occ(mo_energy=None, mo_coeff=None):
+            if mo_energy is None: mo_energy = mf.mo_energy
+            if mo_coeff is None: mo_coeff = mf.mo_coeff
+            if isinstance(mf, rohf.ROHF):
+                mo_coeff = numpy.array([mo_coeff, mo_coeff])
+            mo_occ = numpy.zeros_like(setocc)
+            nocc_a = int(numpy.sum(setocc[0]))
+            nocc_b = int(numpy.sum(setocc[1]))
+            s_a = reduce(numpy.dot, (coef_occ_a.conj().T, mf.get_ovlp(), mo_coeff[0]))
+            s_b = reduce(numpy.dot, (coef_occ_b.conj().T, mf.get_ovlp(), mo_coeff[1]))
+            #choose a subset of mo_coeff, which maximizes <old|now>
+            idx_a = numpy.argsort(numpy.einsum('ij,ij->j', s_a, s_a))[::-1]
+            idx_b = numpy.argsort(numpy.einsum('ij,ij->j', s_b, s_b))[::-1]
+            mo_occ[0,idx_a[:nocc_a]] = 1.
+            mo_occ[1,idx_b[:nocc_b]] = 1.
+
+            log.debug(' New alpha occ pattern: %s', mo_occ[0])
+            log.debug(' New beta occ pattern: %s', mo_occ[1])
+            if isinstance(mf.mo_energy, numpy.ndarray) and mf.mo_energy.ndim == 1:
+                log.debug1(' Current mo_energy(sorted) = %s', mo_energy)
+            else:
+                log.debug1(' Current alpha mo_energy(sorted) = %s', mo_energy[0])
+                log.debug1(' Current beta mo_energy(sorted) = %s', mo_energy[1])
+
+            if (int(numpy.sum(mo_occ[0])) != nocc_a):
+                log.error('mom alpha electron occupation numbers do not match: %d, %d',
+                          nocc_a, int(numpy.sum(mo_occ[0])))
+            if (int(numpy.sum(mo_occ[1])) != nocc_b):
+                log.error('mom beta electron occupation numbers do not match: %d, %d',
+                          nocc_b, int(numpy.sum(mo_occ[1])))
+
+            #output 1-dimension occupation number for restricted open-shell
+            if isinstance(mf, rohf.ROHF): mo_occ = mo_occ[0, :] + mo_occ[1, :]
+            return mo_occ
     else:
-        raise RuntimeError('Cannot support this class of instance %s' % mf)
-    log = logger.Logger(mf.stdout, mf.verbose)
-    def get_occ(mo_energy=None, mo_coeff=None):
-        if mo_energy is None: mo_energy = mf.mo_energy
-        if mo_coeff is None: mo_coeff = mf.mo_coeff
-        if isinstance(mf, rohf.ROHF): mo_coeff = numpy.array([mo_coeff, mo_coeff])
-        mo_occ = numpy.zeros_like(setocc)
-        nocc_a = int(numpy.sum(setocc[0]))
-        nocc_b = int(numpy.sum(setocc[1]))
-        s_a = reduce(numpy.dot, (coef_occ_a.T, mf.get_ovlp(), mo_coeff[0]))
-        s_b = reduce(numpy.dot, (coef_occ_b.T, mf.get_ovlp(), mo_coeff[1]))
-        #choose a subset of mo_coeff, which maximizes <old|now>
-        idx_a = numpy.argsort(numpy.einsum('ij,ij->j', s_a, s_a))[::-1]
-        idx_b = numpy.argsort(numpy.einsum('ij,ij->j', s_b, s_b))[::-1]
-        mo_occ[0][idx_a[:nocc_a]] = 1.
-        mo_occ[1][idx_b[:nocc_b]] = 1.
-
-        log.debug(' New alpha occ pattern: %s', mo_occ[0])
-        log.debug(' New beta occ pattern: %s', mo_occ[1])
-        if isinstance(mf.mo_energy, numpy.ndarray) and mf.mo_energy.ndim == 1:
-            log.debug1(' Current mo_energy(sorted) = %s', mo_energy)
-        else:
-            log.debug1(' Current alpha mo_energy(sorted) = %s', mo_energy[0])
-            log.debug1(' Current beta mo_energy(sorted) = %s', mo_energy[1])
-
-        if (int(numpy.sum(mo_occ[0])) != nocc_a):
-            log.error('mom alpha electron occupation numbers do not match: %d, %d',
-                      nocc_a, int(numpy.sum(mo_occ[0])))
-        if (int(numpy.sum(mo_occ[1])) != nocc_b):
-            log.error('mom beta electron occupation numbers do not match: %d, %d',
-                      nocc_b, int(numpy.sum(mo_occ[1])))
-
-        #output 1-dimension occupation number for restricted open-shell
-        if isinstance(mf, rohf.ROHF): mo_occ = mo_occ[0, :] + mo_occ[1, :]
-        return mo_occ
+        def get_occ(mo_energy=None, mo_coeff=None):
+            if mo_energy is None: mo_energy = mf.mo_energy
+            if mo_coeff is None: mo_coeff = mf.mo_coeff
+            mo_occ = numpy.zeros_like(setocc)
+            nocc = int(setocc.sum())
+            s = occorb[:,setocc>0].conj().T.dot(mf.get_ovlp()).dot(mo_coeff)
+            #choose a subset of mo_coeff, which maximizes <old|now>
+            idx = numpy.argsort(numpy.einsum('ij,ij->j', s, s))[::-1]
+            mo_occ[idx[:nocc]] = 1.
+            return mo_occ
+
     mf.get_occ = get_occ
     return mf
+
 mom_occ = mom_occ_
 
 def project_mo_nr2nr(mol1, mo1, mol2):