#!/usr/bin/env python # Copyright 2014-2019 The PySCF Developers. All Rights Reserved. # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. # You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # # Author: Qiming Sun # import copy import numpy from pyscf import gto from pyscf.lib import logger from pyscf.lib import param from pyscf.data import elements from pyscf.scf import hf, rohf def get_atm_nrhf(mol, atomic_configuration=elements.NRSRHF_CONFIGURATION): elements = set([a[0] for a in mol._atom]) logger.info(mol, 'Spherically averaged atomic HF for %s', elements) atm_template = copy.copy(mol) atm_template.charge = 0 atm_template.symmetry = False # TODO: enable SO3 symmetry here atm_template.atom = atm_template._atom = [] atm_template.cart = False # AtomSphAverageRHF does not support cartensian basis atm_scf_result = {} for ia, a in enumerate(mol._atom): element = a[0] if element in atm_scf_result: continue atm = atm_template atm._atom = [a] atm._atm = mol._atm[ia:ia+1] atm._bas = mol._bas[mol._bas[:,0] == ia].copy() atm._ecpbas = mol._ecpbas[mol._ecpbas[:,0] == ia].copy() # Point to the only atom atm._bas[:,0] = 0 atm._ecpbas[:,0] = 0 if element in mol._pseudo: atm._pseudo = {element: mol._pseudo.get(element)} raise NotImplementedError atm.spin = atm.nelectron % 2 nao = atm.nao # nao == 0 for the case that no basis was assigned to an atom if nao == 0 or atm.nelectron == 0: # GHOST mo_occ = mo_energy = numpy.zeros(nao) mo_coeff = numpy.zeros((nao,nao)) atm_scf_result[element] = (0, mo_energy, mo_coeff, mo_occ) else: if atm.nelectron == 1: atm_hf = AtomHF1e(atm) else: atm_hf = AtomSphAverageRHF(atm) atm_hf.atomic_configuration = atomic_configuration atm_hf.verbose = mol.verbose atm_hf.run() atm_scf_result[element] = (atm_hf.e_tot, atm_hf.mo_energy, atm_hf.mo_coeff, atm_hf.mo_occ) return atm_scf_result class AtomSphAverageRHF(hf.RHF): def __init__(self, mol): self._eri = None self.atomic_configuration = elements.NRSRHF_CONFIGURATION hf.SCF.__init__(self, mol) # The default initial guess minao does not have super-heavy elements if mol.atom_charge(0) > 96: self.init_guess = '1e' def check_sanity(self): pass def dump_flags(self, verbose=None): log = logger.new_logger(self, verbose) hf.RHF.dump_flags(self, log) log.info('atom = %s', self.mol.atom_symbol(0)) def eig(self, f, s): mol = self.mol ao_ang = _angular_momentum_for_each_ao(mol) nao = mol.nao mo_coeff = [] mo_energy = [] for l in range(param.L_MAX): degen = 2 * l + 1 idx = numpy.where(ao_ang == l)[0] nao_l = len(idx) if nao_l > 0: nsh = nao_l // degen f_l = f[idx[:,None],idx].reshape(nsh, degen, nsh, degen) s_l = s[idx[:,None],idx].reshape(nsh, degen, nsh, degen) # Average over angular parts f_l = numpy.einsum('piqi->pq', f_l) / degen s_l = numpy.einsum('piqi->pq', s_l) / degen e, c = self._eigh(f_l, s_l) for i, ei in enumerate(e): logger.debug1(self, 'l = %d e_%d = %.9g', l, i, ei) mo_energy.append(numpy.repeat(e, degen)) mo = numpy.zeros((nao, nsh, degen)) for i in range(degen): mo[idx[i::degen],:,i] = c mo_coeff.append(mo.reshape(nao, nao_l)) return numpy.hstack(mo_energy), numpy.hstack(mo_coeff) def get_occ(self, mo_energy=None, mo_coeff=None): '''spherically averaged fractional occupancy''' mol = self.mol symb = mol.atom_symbol(0) nelec_ecp = mol.atom_nelec_core(0) coreshl = gto.ecp.core_configuration(nelec_ecp) occ = [] for l in range(param.L_MAX): n2occ, frac = frac_occ(symb, l, self.atomic_configuration) degen = 2 * l + 1 idx = mol._bas[:,gto.ANG_OF] == l nbas_l = mol._bas[idx,gto.NCTR_OF].sum() if l < 4: n2occ -= coreshl[l] assert n2occ <= nbas_l logger.debug1(self, 'l = %d occ = %d + %.4g', l, n2occ, frac) occ_l = numpy.zeros(nbas_l) occ_l[:n2occ] = 2 if frac > 0: occ_l[n2occ] = frac occ.append(numpy.repeat(occ_l, degen)) else: occ.append(numpy.zeros(nbas_l * degen)) return numpy.hstack(occ) def get_grad(self, mo_coeff, mo_occ, fock=None): return 0 def scf(self, *args, **kwargs): kwargs['dump_chk'] = False return hf.RHF.scf(self, *args, **kwargs) def _finalize(self): if self.converged: logger.info(self, 'Atomic HF for atom %s converged. SCF energy = %.15g', self.mol.atom_symbol(0), self.e_tot) else: logger.info(self, 'Atomic HF for atom %s not converged. SCF energy = %.15g', self.mol.atom_symbol(0), self.e_tot) return self AtomSphericAverageRHF = AtomSphAverageRHF class AtomHF1e(rohf.HF1e, AtomSphAverageRHF): eig = AtomSphAverageRHF.eig def frac_occ(symb, l, atomic_configuration=elements.NRSRHF_CONFIGURATION): nuc = gto.charge(symb) if l < 4 and atomic_configuration[nuc][l] > 0: ne = atomic_configuration[nuc][l] nd = (l * 2 + 1) * 2 ndocc = ne.__floordiv__(nd) frac = (float(ne) / nd - ndocc) * 2 else: ndocc = frac = 0 return ndocc, frac def _angular_momentum_for_each_ao(mol): ao_ang = numpy.zeros(mol.nao, dtype=numpy.int) ao_loc = mol.ao_loc_nr() for i in range(mol.nbas): p0, p1 = ao_loc[i], ao_loc[i+1] ao_ang[p0:p1] = mol.bas_angular(i) return ao_ang if __name__ == '__main__': mol = gto.Mole() mol.verbose = 5 mol.output = None mol.atom = [["N", (0. , 0., .5)], ["N", (0. , 0.,-.5)] ] mol.basis = {"N": '6-31g'} mol.build() print(get_atm_nrhf(mol))