1#!/usr/bin/env python
2# Copyright 2014-2019 The PySCF Developers. All Rights Reserved.
3#
4# Licensed under the Apache License, Version 2.0 (the "License");
5# you may not use this file except in compliance with the License.
6# You may obtain a copy of the License at
7#
8#     http://www.apache.org/licenses/LICENSE-2.0
9#
10# Unless required by applicable law or agreed to in writing, software
11# distributed under the License is distributed on an "AS IS" BASIS,
12# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13# See the License for the specific language governing permissions and
14# limitations under the License.
15#
16# Author: Qiming Sun <osirpt.sun@gmail.com>
17#
18
19'''
20Non-relativistic generalized Hartree-Fock with point group symmetry.
21'''
22
23from functools import reduce
24import numpy
25import scipy.linalg
26from pyscf import lib
27from pyscf import symm
28from pyscf.lib import logger
29from pyscf.scf import hf_symm
30from pyscf.scf import ghf
31from pyscf.scf import chkfile
32from pyscf import __config__
33
34WITH_META_LOWDIN = getattr(__config__, 'scf_analyze_with_meta_lowdin', True)
35MO_BASE = getattr(__config__, 'MO_BASE', 1)
36
37
38def analyze(mf, verbose=logger.DEBUG, with_meta_lowdin=WITH_META_LOWDIN,
39            **kwargs):
40    mol = mf.mol
41    if not mol.symmetry:
42        return ghf.analyze(mf, verbose, **kwargs)
43
44    mo_energy = mf.mo_energy
45    mo_occ = mf.mo_occ
46    mo_coeff = mf.mo_coeff
47    ovlp_ao = mf.get_ovlp()
48    log = logger.new_logger(mf, verbose)
49    if log.verbose >= logger.NOTE:
50        nirrep = len(mol.irrep_id)
51        orbsym = mf.get_orbsym(mo_coeff, ovlp_ao)
52        wfnsym = 0
53        noccs = [sum(orbsym[mo_occ>0]==ir) for ir in mol.irrep_id]
54        log.note('total symmetry = %s', symm.irrep_id2name(mol.groupname, wfnsym))
55        log.note('occupancy for each irrep:  ' + (' %4s'*nirrep), *mol.irrep_name)
56        log.note('double occ                 ' + (' %4d'*nirrep), *noccs)
57        log.note('**** MO energy ****')
58        irname_full = {}
59        for k,ir in enumerate(mol.irrep_id):
60            irname_full[ir] = mol.irrep_name[k]
61        irorbcnt = {}
62        for k, j in enumerate(orbsym):
63            if j in irorbcnt:
64                irorbcnt[j] += 1
65            else:
66                irorbcnt[j] = 1
67            log.note('MO #%d (%s #%d), energy= %.15g occ= %g',
68                     k+MO_BASE, irname_full[j], irorbcnt[j], mo_energy[k],
69                     mo_occ[k])
70
71    dm = mf.make_rdm1(mo_coeff, mo_occ)
72    dip = mf.dip_moment(mol, dm, verbose=log)
73    if with_meta_lowdin:
74        pop_and_chg = mf.mulliken_meta(mol, dm, s=ovlp_ao, verbose=log)
75    else:
76        pop_and_chg = mf.mulliken_pop(mol, dm, s=ovlp_ao, verbose=log)
77    return pop_and_chg, dip
78
79def canonicalize(mf, mo_coeff, mo_occ, fock=None):
80    '''Canonicalization diagonalizes the UHF Fock matrix in occupied, virtual
81    subspaces separatedly (without change occupancy).
82    '''
83    mol = mf.mol
84    if not mol.symmetry:
85        return ghf.canonicalize(mf, mo_coeff, mo_occ, fock)
86
87    if getattr(mo_coeff, 'orbsym', None) is not None:
88        return hf_symm.canonicalize(mf, mo_coeff, mo_occ, fock)
89    else:
90        raise NotImplementedError
91
92class GHF(ghf.GHF):
93    __doc__ = ghf.GHF.__doc__ + '''
94    Attributes for symmetry allowed GHF:
95        irrep_nelec : dict
96            Specify the number of electrons for particular irrep
97            {'ir_name':int, ...}.
98            For the irreps not listed in these dicts, the program will choose the
99            occupancy based on the orbital energies.
100    '''
101    def __init__(self, mol):
102        ghf.GHF.__init__(self, mol)
103        # number of electrons for each irreps
104        self.irrep_nelec = {}
105        self._keys = self._keys.union(['irrep_nelec'])
106
107    def dump_flags(self, verbose=None):
108        ghf.GHF.dump_flags(self, verbose)
109        if self.irrep_nelec:
110            logger.info(self, 'irrep_nelec %s', self.irrep_nelec)
111        return self
112
113    def build(self, mol=None):
114        if mol is None: mol = self.mol
115        if mol.symmetry:
116            for irname in self.irrep_nelec:
117                if irname not in mol.irrep_name:
118                    logger.warn(self, 'Molecule does not have irrep %s', irname)
119
120            nelec_fix = self.irrep_nelec.values()
121            if any(isinstance(x, (tuple, list)) for x in nelec_fix):
122                msg =('Number of alpha/beta electrons cannot be assigned '
123                      'separately in GHF.  irrep_nelec = %s' % self.irrep_nelec)
124                raise ValueError(msg)
125            nelec_fix = sum(nelec_fix)
126            float_irname = set(mol.irrep_name) - set(self.irrep_nelec)
127            if nelec_fix > mol.nelectron:
128                msg =('More electrons defined by irrep_nelec than total num electrons. '
129                      'mol.nelectron = %d  irrep_nelec = %s' %
130                      (mol.nelectron, self.irrep_nelec))
131                raise ValueError(msg)
132            else:
133                logger.info(mol, 'Freeze %d electrons in irreps %s',
134                            nelec_fix, self.irrep_nelec.keys())
135
136            if len(float_irname) == 0 and nelec_fix != mol.nelectron:
137                msg =('Num electrons defined by irrep_nelec != total num electrons. '
138                      'mol.nelectron = %d  irrep_nelec = %s' %
139                      (mol.nelectron, self.irrep_nelec))
140                raise ValueError(msg)
141            else:
142                logger.info(mol, '    %d free electrons in irreps %s',
143                            mol.nelectron-nelec_fix, ' '.join(float_irname))
144        return ghf.GHF.build(self, mol)
145
146    def eig(self, h, s):
147        mol = self.mol
148        if not mol.symmetry:
149            return self._eigh(h, s)
150
151        nirrep = len(mol.symm_orb)
152        symm_orb = [scipy.linalg.block_diag(c, c) for c in mol.symm_orb]
153        s = [reduce(numpy.dot, (c.T,s,c)) for c in symm_orb]
154        h = [reduce(numpy.dot, (c.T,h,c)) for c in symm_orb]
155        cs = []
156        es = []
157        orbsym = []
158        for ir in range(nirrep):
159            e, c = self._eigh(h[ir], s[ir])
160            cs.append(c)
161            es.append(e)
162            orbsym.append([mol.irrep_id[ir]] * e.size)
163        e = numpy.hstack(es)
164        c = hf_symm.so2ao_mo_coeff(symm_orb, cs)
165        c = lib.tag_array(c, orbsym=numpy.hstack(orbsym))
166        return e, c
167
168    def get_grad(self, mo_coeff, mo_occ, fock=None):
169        g = ghf.GHF.get_grad(self, mo_coeff, mo_occ, fock)
170        if self.mol.symmetry:
171            occidx = mo_occ > 0
172            viridx = ~occidx
173            orbsym = self.get_orbsym(mo_coeff, self.get_ovlp())
174            sym_forbid = orbsym[viridx].reshape(-1,1) != orbsym[occidx]
175            g[sym_forbid.ravel()] = 0
176        return g
177
178    def get_occ(self, mo_energy=None, mo_coeff=None):
179        ''' We assumed mo_energy are grouped by symmetry irreps, (see function
180        self.eig). The orbitals are sorted after SCF.
181        '''
182        if mo_energy is None: mo_energy = self.mo_energy
183        mol = self.mol
184        if not mol.symmetry:
185            return ghf.GHF.get_occ(self, mo_energy, mo_coeff)
186
187        orbsym = self.get_orbsym(mo_coeff, self.get_ovlp())
188        mo_occ = numpy.zeros_like(mo_energy)
189        rest_idx = numpy.ones(mo_occ.size, dtype=bool)
190        nelec_fix = 0
191        for i, ir in enumerate(mol.irrep_id):
192            irname = mol.irrep_name[i]
193            ir_idx = numpy.where(orbsym == ir)[0]
194            if irname in self.irrep_nelec:
195                n = self.irrep_nelec[irname]
196                occ_sort = numpy.argsort(mo_energy[ir_idx].round(9), kind='mergesort')
197                occ_idx  = ir_idx[occ_sort[:n]]
198                mo_occ[occ_idx] = 1
199                nelec_fix += n
200                rest_idx[ir_idx] = False
201        nelec_float = mol.nelectron - nelec_fix
202        assert(nelec_float >= 0)
203        if nelec_float > 0:
204            rest_idx = numpy.where(rest_idx)[0]
205            occ_sort = numpy.argsort(mo_energy[rest_idx].round(9), kind='mergesort')
206            occ_idx  = rest_idx[occ_sort[:nelec_float]]
207            mo_occ[occ_idx] = 1
208
209        vir_idx = (mo_occ==0)
210        if self.verbose >= logger.INFO and numpy.count_nonzero(vir_idx) > 0:
211            ehomo = max(mo_energy[~vir_idx])
212            elumo = min(mo_energy[ vir_idx])
213            noccs = []
214            for i, ir in enumerate(mol.irrep_id):
215                irname = mol.irrep_name[i]
216                ir_idx = (orbsym == ir)
217
218                noccs.append(int(mo_occ[ir_idx].sum()))
219                if ehomo in mo_energy[ir_idx]:
220                    irhomo = irname
221                if elumo in mo_energy[ir_idx]:
222                    irlumo = irname
223            logger.info(self, 'HOMO (%s) = %.15g  LUMO (%s) = %.15g',
224                        irhomo, ehomo, irlumo, elumo)
225
226            logger.debug(self, 'irrep_nelec = %s', noccs)
227            hf_symm._dump_mo_energy(mol, mo_energy, mo_occ, ehomo, elumo, orbsym,
228                                    verbose=self.verbose)
229
230            if mo_coeff is not None and self.verbose >= logger.DEBUG:
231                ss, s = self.spin_square(mo_coeff[:,mo_occ>0], self.get_ovlp())
232                logger.debug(self, 'multiplicity <S^2> = %.8g  2S+1 = %.8g', ss, s)
233        return mo_occ
234
235    def _finalize(self):
236        ghf.GHF._finalize(self)
237
238        # Using mergesort because it is stable. We don't want to change the
239        # ordering of the symmetry labels when two orbitals are degenerated.
240        o_sort = numpy.argsort(self.mo_energy[self.mo_occ> 0].round(9), kind='mergesort')
241        v_sort = numpy.argsort(self.mo_energy[self.mo_occ==0].round(9), kind='mergesort')
242        orbsym = self.get_orbsym(self.mo_coeff, self.get_ovlp())
243        self.mo_energy = numpy.hstack((self.mo_energy[self.mo_occ> 0][o_sort],
244                                       self.mo_energy[self.mo_occ==0][v_sort]))
245        self.mo_coeff = numpy.hstack((self.mo_coeff[:,self.mo_occ> 0].take(o_sort, axis=1),
246                                      self.mo_coeff[:,self.mo_occ==0].take(v_sort, axis=1)))
247        orbsym = numpy.hstack((orbsym[self.mo_occ> 0][o_sort],
248                               orbsym[self.mo_occ==0][v_sort]))
249        self.mo_coeff = lib.tag_array(self.mo_coeff, orbsym=orbsym)
250        nocc = len(o_sort)
251        self.mo_occ[:nocc] = 1
252        self.mo_occ[nocc:] = 0
253        if self.chkfile:
254            chkfile.dump_scf(self.mol, self.chkfile, self.e_tot, self.mo_energy,
255                             self.mo_coeff, self.mo_occ, overwrite_mol=False)
256        return self
257
258    def analyze(self, verbose=None, **kwargs):
259        if verbose is None: verbose = self.verbose
260        return analyze(self, verbose, **kwargs)
261
262    @lib.with_doc(hf_symm.get_irrep_nelec.__doc__)
263    def get_irrep_nelec(self, mol=None, mo_coeff=None, mo_occ=None, s=None):
264        if mol is None: mol = self.mol
265        if mo_occ is None: mo_occ = self.mo_occ
266        if mo_coeff is None: mo_coeff = self.mo_coeff
267        if s is None: s = self.get_ovlp()
268        return hf_symm.get_irrep_nelec(mol, mo_coeff, mo_occ, s)
269
270    canonicalize = canonicalize
271
272    def get_orbsym(self, mo_coeff=None, s=None):
273        if mo_coeff is None:
274            mo_coeff = self.mo_coeff
275        if s is None:
276            s = self.get_ovlp()
277        return numpy.asarray(get_orbsym(self.mol, mo_coeff, s))
278    orbsym = property(get_orbsym)
279
280def get_orbsym(mol, mo_coeff, s=None, check=False):
281    if mo_coeff is None:
282        orbsym = numpy.hstack([[ir] * mol.symm_orb[i].shape[1]
283                               for i, ir in enumerate(mol.irrep_id)])
284    elif getattr(mo_coeff, 'orbsym', None) is not None:
285        orbsym = mo_coeff.orbsym
286    else:
287        nao = mo_coeff.shape[0] // 2
288        if isinstance(s, numpy.ndarray):
289            assert(s.size == nao**2 or numpy.allclose(s[:nao,:nao], s[nao:,nao:]))
290            s = s[:nao,:nao]
291        mo_a = mo_coeff[:nao].copy()
292        mo_b = mo_coeff[nao:]
293        zero_alpha_idx = numpy.linalg.norm(mo_a, axis=0) < 1e-7
294        mo_a[:,zero_alpha_idx] = mo_b[:,zero_alpha_idx]
295        orbsym = symm.label_orb_symm(mol, mol.irrep_id, mol.symm_orb,
296                                     mo_a, s, check)
297    return numpy.asarray(orbsym)
298
299
300if __name__ == '__main__':
301    from pyscf import gto
302    mol = gto.Mole()
303    mol.build(
304        verbose = 1,
305        output = None,
306        atom = [['O', (0.,0.,0.)],
307                ['O', (0.,0.,1.)], ],
308        basis = {'O': 'ccpvdz'},
309        symmetry = True,
310        charge = -1,
311        spin = 1
312    )
313
314    method = GHF(mol)
315    method.verbose = 5
316    method.irrep_nelec['A1u'] = 1
317    energy = method.kernel()
318    print(energy - -126.117033823738)
319    method.canonicalize(method.mo_coeff, method.mo_occ)
320    method.analyze()
321