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
19import copy
20import numpy
21from pyscf import gto
22from pyscf.lib import logger
23from pyscf.lib import param
24from pyscf.data import elements
25from pyscf.scf import hf, rohf
26
27
28def get_atm_nrhf(mol, atomic_configuration=elements.NRSRHF_CONFIGURATION):
29    elements = set([a[0] for a in mol._atom])
30    logger.info(mol, 'Spherically averaged atomic HF for %s', elements)
31
32    atm_template = copy.copy(mol)
33    atm_template.charge = 0
34    atm_template.symmetry = False  # TODO: enable SO3 symmetry here
35    atm_template.atom = atm_template._atom = []
36    atm_template.cart = False  # AtomSphAverageRHF does not support cartensian basis
37
38    atm_scf_result = {}
39    for ia, a in enumerate(mol._atom):
40        element = a[0]
41        if element in atm_scf_result:
42            continue
43
44        atm = atm_template
45        atm._atom = [a]
46        atm._atm = mol._atm[ia:ia+1]
47        atm._bas = mol._bas[mol._bas[:,0] == ia].copy()
48        atm._ecpbas = mol._ecpbas[mol._ecpbas[:,0] == ia].copy()
49        # Point to the only atom
50        atm._bas[:,0] = 0
51        atm._ecpbas[:,0] = 0
52        if element in mol._pseudo:
53            atm._pseudo = {element: mol._pseudo.get(element)}
54            raise NotImplementedError
55        atm.spin = atm.nelectron % 2
56
57        nao = atm.nao
58        # nao == 0 for the case that no basis was assigned to an atom
59        if nao == 0 or atm.nelectron == 0:  # GHOST
60            mo_occ = mo_energy = numpy.zeros(nao)
61            mo_coeff = numpy.zeros((nao,nao))
62            atm_scf_result[element] = (0, mo_energy, mo_coeff, mo_occ)
63        else:
64            if atm.nelectron == 1:
65                atm_hf = AtomHF1e(atm)
66            else:
67                atm_hf = AtomSphAverageRHF(atm)
68                atm_hf.atomic_configuration = atomic_configuration
69
70            atm_hf.verbose = mol.verbose
71            atm_hf.run()
72            atm_scf_result[element] = (atm_hf.e_tot, atm_hf.mo_energy,
73                                       atm_hf.mo_coeff, atm_hf.mo_occ)
74    return atm_scf_result
75
76
77class AtomSphAverageRHF(hf.RHF):
78    def __init__(self, mol):
79        self._eri = None
80        self.atomic_configuration = elements.NRSRHF_CONFIGURATION
81        hf.SCF.__init__(self, mol)
82
83        # The default initial guess minao does not have super-heavy elements
84        if mol.atom_charge(0) > 96:
85            self.init_guess = '1e'
86
87    def check_sanity(self):
88        pass
89
90    def dump_flags(self, verbose=None):
91        log = logger.new_logger(self, verbose)
92        hf.RHF.dump_flags(self, log)
93        log.info('atom = %s', self.mol.atom_symbol(0))
94
95    def eig(self, f, s):
96        mol = self.mol
97        ao_ang = _angular_momentum_for_each_ao(mol)
98
99        nao = mol.nao
100        mo_coeff = []
101        mo_energy = []
102
103        for l in range(param.L_MAX):
104            degen = 2 * l + 1
105            idx = numpy.where(ao_ang == l)[0]
106            nao_l = len(idx)
107
108            if nao_l > 0:
109                nsh = nao_l // degen
110                f_l = f[idx[:,None],idx].reshape(nsh, degen, nsh, degen)
111                s_l = s[idx[:,None],idx].reshape(nsh, degen, nsh, degen)
112                # Average over angular parts
113                f_l = numpy.einsum('piqi->pq', f_l) / degen
114                s_l = numpy.einsum('piqi->pq', s_l) / degen
115
116                e, c = self._eigh(f_l, s_l)
117                for i, ei in enumerate(e):
118                    logger.debug1(self, 'l = %d  e_%d = %.9g', l, i, ei)
119                mo_energy.append(numpy.repeat(e, degen))
120
121                mo = numpy.zeros((nao, nsh, degen))
122                for i in range(degen):
123                    mo[idx[i::degen],:,i] = c
124                mo_coeff.append(mo.reshape(nao, nao_l))
125
126        return numpy.hstack(mo_energy), numpy.hstack(mo_coeff)
127
128    def get_occ(self, mo_energy=None, mo_coeff=None):
129        '''spherically averaged fractional occupancy'''
130        mol = self.mol
131        symb = mol.atom_symbol(0)
132
133        nelec_ecp = mol.atom_nelec_core(0)
134        coreshl = gto.ecp.core_configuration(nelec_ecp)
135
136        occ = []
137        for l in range(param.L_MAX):
138            n2occ, frac = frac_occ(symb, l, self.atomic_configuration)
139            degen = 2 * l + 1
140            idx = mol._bas[:,gto.ANG_OF] == l
141            nbas_l = mol._bas[idx,gto.NCTR_OF].sum()
142            if l < 4:
143                n2occ -= coreshl[l]
144                assert n2occ <= nbas_l
145
146                logger.debug1(self, 'l = %d  occ = %d + %.4g', l, n2occ, frac)
147
148                occ_l = numpy.zeros(nbas_l)
149                occ_l[:n2occ] = 2
150                if frac > 0:
151                    occ_l[n2occ] = frac
152                occ.append(numpy.repeat(occ_l, degen))
153            else:
154                occ.append(numpy.zeros(nbas_l * degen))
155
156        return numpy.hstack(occ)
157
158    def get_grad(self, mo_coeff, mo_occ, fock=None):
159        return 0
160
161    def scf(self, *args, **kwargs):
162        kwargs['dump_chk'] = False
163        return hf.RHF.scf(self, *args, **kwargs)
164
165    def _finalize(self):
166        if self.converged:
167            logger.info(self, 'Atomic HF for atom  %s  converged. SCF energy = %.15g',
168                        self.mol.atom_symbol(0), self.e_tot)
169        else:
170            logger.info(self, 'Atomic HF for atom  %s  not converged. SCF energy = %.15g',
171                        self.mol.atom_symbol(0), self.e_tot)
172        return self
173
174AtomSphericAverageRHF = AtomSphAverageRHF
175
176
177class AtomHF1e(rohf.HF1e, AtomSphAverageRHF):
178    eig = AtomSphAverageRHF.eig
179
180
181def frac_occ(symb, l, atomic_configuration=elements.NRSRHF_CONFIGURATION):
182    nuc = gto.charge(symb)
183    if l < 4 and atomic_configuration[nuc][l] > 0:
184        ne = atomic_configuration[nuc][l]
185        nd = (l * 2 + 1) * 2
186        ndocc = ne.__floordiv__(nd)
187        frac = (float(ne) / nd - ndocc) * 2
188    else:
189        ndocc = frac = 0
190    return ndocc, frac
191
192
193def _angular_momentum_for_each_ao(mol):
194    ao_ang = numpy.zeros(mol.nao, dtype=numpy.int)
195    ao_loc = mol.ao_loc_nr()
196    for i in range(mol.nbas):
197        p0, p1 = ao_loc[i], ao_loc[i+1]
198        ao_ang[p0:p1] = mol.bas_angular(i)
199    return ao_ang
200
201
202if __name__ == '__main__':
203    mol = gto.Mole()
204    mol.verbose = 5
205    mol.output = None
206
207    mol.atom = [["N", (0. , 0., .5)],
208                ["N", (0. , 0.,-.5)] ]
209
210    mol.basis = {"N": '6-31g'}
211    mol.build()
212    print(get_atm_nrhf(mol))
213