Update PM localizatin schemes (issue pyscf#563)

sichaox · May 21, 2020 · 3fe8ce0 · 3fe8ce0
1 parent b0e83b6
commit 3fe8ce0
Show file tree

Hide file tree

Showing 13 changed files with 174 additions and 62 deletions.
diff --git a/.flake8 b/.flake8
@@ -2,7 +2,7 @@
 # https://flake8.pycqa.org/en/2.5.5/warnings.html#error-codes
 ignore = E127, E129, E201, E202, E203, E225, E226, E231, E241, E261, E262, \
          E265, E266, E301, E302, E303, E305, E401, E701, W391, W504, C901, \
-         F401, E228, E126, W503, E731, E211, E221
+         F401, E228, E126, W503, E731, E211, E221, E222, E741
 
 exclude = test, .git, __pycache__, build, dist, __init__.py .eggs, *.egg
 max-line-length = 120
diff --git a/examples/local_orb/02-pop_with_iao.py b/examples/local_orb/02-pop_with_iao.py
@@ -3,10 +3,10 @@
 '''
 Mulliken population analysis with IAO orbitals
 '''
+
 import numpy
 from functools import reduce
 from pyscf import gto, scf, lo
-from functools import reduce
 
 x = .63
 mol = gto.M(atom=[['C', (0, 0, 0)],

diff --git a/examples/local_orb/07-pipek_mezey.py b/examples/local_orb/07-pipek_mezey.py
@@ -0,0 +1,44 @@
+#!/usr/bin/env python
+
+'''
+Various Pipek-Mezey localization schemes
+
+See more discussions in paper  JCTC, 10, 642 (2014); DOI:10.1021/ct401016x
+'''
+
+import pyscf
+
+x = .63
+mol = pyscf.M(atom=[['C', (0, 0, 0)],
+                    ['H', (x ,  x,  x)],
+                    ['H', (-x, -x,  x)],
+                    ['H', (-x,  x, -x)],
+                    ['H', ( x, -x, -x)]],
+              basis='ccpvtz')
+
+mf = mol.RHF().run()
+
+orbitals = mf.mo_coeff[:,mf.mo_occ>0]
+pm = pyscf.lo.PM(mol, orbitals, mf)
+def print_coeff(local_orb):
+    import sys
+    idx = mol.search_ao_label(['C 1s', 'C 2s', 'C 2p', 'H 1s'])
+    ao_labels = mol.ao_labels()
+    labels = [ao_labels[i] for i in idx]
+    pyscf.tools.dump_mat.dump_rec(sys.stdout, local_orb[idx], labels)
+
+print('---')
+pm.pop_method = 'mulliken'
+print_coeff(pm.kernel())
+
+print('---')
+pm.pop_method = 'meta-lowdin'
+print_coeff(pm.kernel())
+
+print('---')
+pm.pop_method = 'iao'
+print_coeff(pm.kernel())
+
+print('---')
+pm.pop_method = 'becke'
+print_coeff(pm.kernel())
diff --git a/pyscf/dft/gen_grid.py b/pyscf/dft/gen_grid.py
@@ -275,12 +275,16 @@ def gen_atomic_grids(mol, atom_grid={}, radi_method=radi.gauss_chebyshev,
     return atom_grids_tab
 
 
-def gen_partition(mol, atom_grids_tab,
+def get_partition(mol, atom_grids_tab,
                   radii_adjust=None, atomic_radii=radi.BRAGG_RADII,
-                  becke_scheme=original_becke):
+                  becke_scheme=original_becke, concat=True):
     '''Generate the mesh grid coordinates and weights for DFT numerical integration.
     We can change radii_adjust, becke_scheme functions to generate different meshgrid.
 
+    Kwargs:
+        concat: bool
+            Whether to concatenate grids and weights in return
+
     Returns:
         grid_coord and grid_weight arrays.  grid_coord array has shape (N,3);
         weight 1D array has N elements.
@@ -339,7 +343,12 @@ def gen_grid_partition(coords):
         weights = vol * pbecke[ia] * (1./pbecke.sum(axis=0))
         coords_all.append(coords)
         weights_all.append(weights)
-    return numpy.vstack(coords_all), numpy.hstack(weights_all)
+
+    if concat:
+        coords_all = numpy.vstack(coords_all)
+        weights_all = numpy.hstack(weights_all)
+    return coords_all, weights_all
+gen_partition = get_partition
 
 def make_mask(mol, coords, relativity=0, shls_slice=None, verbose=None):
     '''Mask to indicate whether a shell is zero on grid
@@ -525,7 +534,7 @@ def build(self, mol=None, with_non0tab=False, **kwargs):
                                                self.radi_method,
                                                self.level, self.prune, **kwargs)
         self.coords, self.weights = \
-                self.gen_partition(mol, atom_grids_tab,
+                self.get_partition(mol, atom_grids_tab,
                                    self.radii_adjust, self.atomic_radii,
                                    self.becke_scheme)
         if with_non0tab:
@@ -551,20 +560,22 @@ def reset(self, mol=None):
     @lib.with_doc(gen_atomic_grids.__doc__)
     def gen_atomic_grids(self, mol, atom_grid=None, radi_method=None,
                          level=None, prune=None, **kwargs):
-        ''' See gen_grid.gen_atomic_grids function'''
         if atom_grid is None: atom_grid = self.atom_grid
         if radi_method is None: radi_method = self.radi_method
         if level is None: level = self.level
         if prune is None: prune = self.prune
         return gen_atomic_grids(mol, atom_grid, self.radi_method, level, prune, **kwargs)
 
-    @lib.with_doc(gen_partition.__doc__)
-    def gen_partition(self, mol, atom_grids_tab,
+    @lib.with_doc(get_partition.__doc__)
+    def get_partition(self, mol, atom_grids_tab=None,
                       radii_adjust=None, atomic_radii=radi.BRAGG_RADII,
-                      becke_scheme=original_becke):
-        ''' See gen_grid.gen_partition function'''
-        return gen_partition(mol, atom_grids_tab, radii_adjust, atomic_radii,
-                             becke_scheme)
+                      becke_scheme=original_becke, concat=True):
+        if atom_grids_tab is None:
+            atom_grids_tab = self.gen_atomic_grids(mol)
+        return get_partition(mol, atom_grids_tab, radii_adjust, atomic_radii,
+                             becke_scheme, concat=concat)
+
+    gen_partition = get_partition
 
     @lib.with_doc(make_mask.__doc__)
     def make_mask(self, mol=None, coords=None, relativity=0, shls_slice=None,

diff --git a/pyscf/lo/boys.py b/pyscf/lo/boys.py
@@ -248,6 +248,7 @@ def gen_g_hop(self, u):
         #:h = h[idx][:,idx[0],idx[1]]
 
         g0 = g0 + g0.conj().T
+
         def h_op(x):
             x = self.unpack_uniq_var(x)
             norb = x.shape[0]

diff --git a/pyscf/lo/edmiston.py b/pyscf/lo/edmiston.py
@@ -51,6 +51,7 @@ def gen_g_hop(self, u):
         h_diag = -self.pack_uniq_var(h_diag) * 2
 
         g0 = g0 + g0.T
+
         def h_op(x):
             x = self.unpack_uniq_var(x)
             hx = numpy.einsum('iq,qp->pi', g0, x)

diff --git a/pyscf/lo/iao.py b/pyscf/lo/iao.py
@@ -32,7 +32,8 @@
 from pyscf.lo.orth import vec_lowdin
 
 # Alternately, use ANO for minao
-# orthogonalize iao by orth.lowdin(c.T*mol.intor(ovlp)*c)
+# orthogonalize iao with coefficients obtained by
+#     vec_lowdin(iao_coeff, mol.intor('int1e_ovlp'))
 MINAO = getattr(__config__, 'lo_iao_minao', 'minao')
 
 def iao(mol, orbocc, minao=MINAO, kpts=None):
@@ -67,13 +68,13 @@ def iao(mol, orbocc, minao=MINAO, kpts=None):
         s2 = numpy.asarray(pmol.pbc_intor('int1e_ovlp', hermi=1, kpts=kpts))
         s12 = numpy.asarray(pbcgto.cell.intor_cross('int1e_ovlp', mol, pmol, kpts=kpts))
     else:
-#s1 is the one electron overlap integrals (coulomb integrals)
+        #s1 is the one electron overlap integrals (coulomb integrals)
         s1 = mol.intor_symmetric('int1e_ovlp')
-#s2 is the same as s1 except in minao 
+        #s2 is the same as s1 except in minao
         s2 = pmol.intor_symmetric('int1e_ovlp')
-#overlap integrals of the two molecules 
+        #overlap integrals of the two molecules
         s12 = gto.mole.intor_cross('int1e_ovlp', mol, pmol)
-#transpose of overlap
+
     if len(s1.shape) == 2:
         s21 = s12.conj().T
         s1cd = scipy.linalg.cho_factor(s1)
@@ -86,7 +87,7 @@ def iao(mol, orbocc, minao=MINAO, kpts=None):
         ccs2 = reduce(numpy.dot, (ctild, ctild.conj().T, s1))
         #a is the set of IAOs in the original basis
         a = (p12 + reduce(numpy.dot, (ccs1, ccs2, p12)) * 2
-            - numpy.dot(ccs1, p12) - numpy.dot(ccs2, p12))
+             - numpy.dot(ccs1, p12) - numpy.dot(ccs2, p12))
     else: # k point sampling
         s21 = numpy.swapaxes(s12, -1, -2).conj()
         nkpts = len(kpts)
@@ -103,7 +104,7 @@ def iao(mol, orbocc, minao=MINAO, kpts=None):
             ccs2_k = reduce(numpy.dot, (ctild_k, ctild_k.conj().T, s1[k]))
             #a is the set of IAOs in the original basis
             a[k] = (p12_k + reduce(numpy.dot, (ccs1_k, ccs2_k, p12_k)) * 2
-                - numpy.dot(ccs1_k, p12_k) - numpy.dot(ccs2_k, p12_k))
+                    - numpy.dot(ccs1_k, p12_k) - numpy.dot(ccs2_k, p12_k))
     return a
 
 def reference_mol(mol, minao=MINAO):

diff --git a/pyscf/lo/ibo.py b/pyscf/lo/ibo.py
@@ -36,7 +36,7 @@
 
 def ibo(mol, orbocc, locmethod='IBO', iaos=None, s=None, exponent=4, grad_tol=1e-8, max_iter=200, verbose=logger.NOTE):
     '''Intrinsic Bonding Orbitals
-    
+
     This function serves as a wrapper to the underlying localization functions
     ibo_loc and PipekMezey to create IBOs.
 
@@ -58,7 +58,7 @@ def ibo(mol, orbocc, locmethod='IBO', iaos=None, s=None, exponent=4, grad_tol=1e
     Returns:
         IBOs in the basis defined in mol object.
     '''
-    
+
     if s is None:
         if getattr(mol, 'pbc_intor', None):  # whether mol object is a cell
             if isinstance(orbocc, numpy.ndarray) and orbocc.ndim == 2:
@@ -81,7 +81,7 @@ def ibo(mol, orbocc, locmethod='IBO', iaos=None, s=None, exponent=4, grad_tol=1e
     return ibos
 
 def ibo_loc(mol, orbocc, iaos, s, exponent, grad_tol, max_iter,
-        verbose=logger.NOTE):
+            verbose=logger.NOTE):
     '''Intrinsic Bonding Orbitals. [Ref. JCTC, 9, 4834]
 
     This implementation follows Knizia's implementation execept that the
@@ -228,11 +228,15 @@ def PipekMezey(mol, orbocc, iaos, s, exponent):
     >>> loc_orb = pm.kernel()
     '''
 
+    # Note: PM with Lowdin-orth IAOs is implemented in pipek.PM class
+    # TODO: Merge the implemenation here to pipek.PM
+
     MINAO = getattr(__config__, 'lo_iao_minao', 'minao')
     cs = numpy.dot(iaos.T.conj(), s)
     s_iao = numpy.dot(cs, iaos)
     iao_inv = numpy.linalg.solve(s_iao, cs)
     iao_mol = iao.reference_mol(mol, minao=MINAO)
+
     # Define the mulliken population of each atom based on IAO basis.
     # proj[i].trace is the mulliken population of atom i.
     def atomic_pops(mol, mo_coeff, method=None):
@@ -291,6 +295,7 @@ def MakeAtomInfos():
     for At in "Ga Ge As Se Br Kr".split(): nAoX[At] = 18/2+1+5+3
 
     AoLabels = {}
+
     def SetAo(At, AoDecl):
         Labels = AoDecl.split()
         AoLabels[At] = Labels

diff --git a/pyscf/lo/nao.py b/pyscf/lo/nao.py
@@ -161,7 +161,7 @@ def _core_val_ryd_list(mol):
     k = 0
     for ib in range(mol.nbas):
         ia = mol.bas_atom(ib)
-# Avoid calling mol.atom_charge because we should include ECP core electrons here
+        # Avoid calling mol.atom_charge because we should include ECP core electrons here
         nuc = mole.charge(mol.atom_symbol(ia))
         l = mol.bas_angular(ib)
         nc = mol.bas_nctr(ib)
@@ -211,6 +211,7 @@ def set_atom_conf(element, description):
             | ("3s2p","1d") : 3 s, 2 p shells for core and 1 d shells for valence
     '''
     charge = mole.charge(element)
+
     def to_conf(desc):
         desc = desc.replace(' ','').replace('-','').replace('_','').lower()
         if "doublep" in desc:
@@ -223,6 +224,7 @@ def to_conf(desc):
             loc = AOSHELL[charge][1].find('0')
             desc = '1' + AOSHELL[charge][1][loc+1]
         return desc
+
     if isinstance(description, str):
         c_desc, v_desc = AOSHELL[charge][0], to_conf(description)
     else:

diff --git a/pyscf/lo/orth.py b/pyscf/lo/orth.py
@@ -65,7 +65,7 @@ def pre_orth_ao(mol, method=REF_BASIS):
     '''Restore AO characters.  Possible methods include the ANO/MINAO
     projection or fraction-averaged atomic RHF calculation'''
     if method.upper() in ('ANO', 'MINAO'):
-# Use ANO/MINAO basis to define the strongly occupied set
+        # Use ANO/MINAO basis to define the strongly occupied set
         return project_to_atomic_orbitals(mol, method)
     else:
         return pre_orth_ao_atm_scf(mol)
@@ -77,6 +77,7 @@ def project_to_atomic_orbitals(mol, basname):
     from pyscf.scf.addons import project_mo_nr2nr
     from pyscf.scf import atom_hf
     from pyscf.gto.ecp import core_configuration
+
     def search_atm_l(atm, l):
         bas_ang = atm._bas[:,gto.ANG_OF]
         ao_loc = atm.ao_loc_nr()
@@ -162,16 +163,16 @@ def ecp_ano_det_ovlp(atm_ecp, atm_ano, ecpcore):
         atmp._built = True
 
         if symb in nelec_ecp_dic and nelec_ecp_dic[symb] > 0:
+            # If ECP basis has good atomic character, ECP basis can be used in the
+            # localization/population analysis directly. Otherwise project ECP
+            # basis to ANO basis.
             if not PROJECT_ECP_BASIS:
-# If ECP basis has good atomic character, ECP basis can be used in the
-# localization/population analysis directly. Otherwise project ECP basis to
-# ANO basis.
                 continue
 
             ecpcore = core_configuration(nelec_ecp_dic[symb])
-# Comparing to ANO valence basis, to check whether the ECP basis set has
-# reasonable AO-character contraction.  The ANO valence AO should have
-# significant overlap to ECP basis if the ECP basis has AO-character.
+            # Comparing to ANO valence basis, to check whether the ECP basis set has
+            # reasonable AO-character contraction.  The ANO valence AO should have
+            # significant overlap to ECP basis if the ECP basis has AO-character.
             if abs(ecp_ano_det_ovlp(atm, atmp, ecpcore)) > .1:
                 aos[symb] = numpy.diag(1./numpy.sqrt(s0.diagonal()))
                 continue
@@ -198,7 +199,7 @@ def ecp_ano_det_ovlp(atm_ecp, atm_ano, ecpcore):
                 degen = l * 2 + 1
 
             if nbf_atm_l > nbf_ano_l > 0:
-# For angular l, first place the projected ANO, then the rest AOs.
+                # For angular l, first place the projected ANO, then the rest AOs.
                 sdiag = reduce(numpy.dot, (rm_ano[:,idx].T, s0, rm_ano[:,idx])).diagonal()
                 nleft = (nbf_atm_l - nbf_ano_l) // degen
                 shell_average = numpy.einsum('ij->i', sdiag.reshape(-1,degen))
@@ -275,7 +276,6 @@ def orth_ao(mf_or_mol, method=ORTH_METHOD, pre_orth_ao=None, scf_method=None,
             s = mol.intor_symmetric('int1e_ovlp')
 
     if pre_orth_ao is None:
-#        pre_orth_ao = numpy.eye(mol.nao_nr())
         pre_orth_ao = project_to_atomic_orbitals(mol, REF_BASIS)
 
     if method.lower() == 'lowdin':
@@ -284,8 +284,9 @@ def orth_ao(mf_or_mol, method=ORTH_METHOD, pre_orth_ao=None, scf_method=None,
     elif method.lower() == 'nao':
         assert(mf is not None)
         c_orth = nao.nao(mol, mf, s)
-    else: # meta_lowdin: divide ao into core, valence and Rydberg sets,
-          # orthogonalizing within each set
+    else:
+        # meta_lowdin: partition AOs into core, valence and Rydberg sets,
+        # orthogonalizing within each set
         weight = numpy.ones(pre_orth_ao.shape[0])
         c_orth = nao._nao_sub(mol, weight, pre_orth_ao, s)
     # adjust phase