From 00e8a9819c27b162062c8d3a3d168ff443911572 Mon Sep 17 00:00:00 2001 From: Qiming Sun Date: Thu, 7 Sep 2023 18:35:16 -0700 Subject: [PATCH] Refactor smearing --- pyscf/pbc/scf/addons.py | 380 +++++++++++------------------- pyscf/pbc/scf/test/test_addons.py | 44 ++-- pyscf/scf/addons.py | 200 +++++++++++++++- pyscf/scf/test/test_addons.py | 45 ++++ 4 files changed, 399 insertions(+), 270 deletions(-) diff --git a/pyscf/pbc/scf/addons.py b/pyscf/pbc/scf/addons.py index a5f575e125..638f65a160 100644 --- a/pyscf/pbc/scf/addons.py +++ b/pyscf/pbc/scf/addons.py @@ -32,7 +32,7 @@ from pyscf.pbc.tools import k2gamma from pyscf import __config__ -SMEARING_METHOD = getattr(__config__, 'pbc_scf_addons_smearing_method', 'fermi') +SMEARING_METHOD = mol_addons.SMEARING_METHOD def project_mo_nr2nr(cell1, mo1, cell2, kpts=None): @@ -69,251 +69,167 @@ def _k2k_projection(kpts1, kpts2, Ls): c = expRk1.T.dot(expRk2) * weight return (c*c.conj()).real.copy() -def smearing_(mf, sigma=None, method=SMEARING_METHOD, mu0=None, fix_spin=False): - '''Fermi-Dirac or Gaussian smearing''' - from pyscf.scf import uhf, rohf - from pyscf.scf import ghf - from pyscf.pbc.scf import khf - mf_class = mf.__class__ - is_uhf = isinstance(mf, uhf.UHF) - is_ghf = isinstance(mf, ghf.GHF) - is_rhf = (not is_uhf) and (not is_ghf) - is_khf = isinstance(mf, khf.KSCF) - is_rohf = isinstance(mf, rohf.ROHF) - if is_rohf: - is_rhf = False - - if fix_spin and not (is_uhf or is_rohf): - raise KeyError("fix_spin only supports UHF and ROHF.") - if fix_spin and mu0 is not None: - raise KeyError("fix_spin does not support fix mu0") - - def fermi_smearing_occ(m, mo_energy_kpts, sigma): - occ = numpy.zeros_like(mo_energy_kpts) - de = (mo_energy_kpts - m) / sigma - occ[de<40] = 1./(numpy.exp(de[de<40])+1.) - return occ - def gaussian_smearing_occ(m, mo_energy_kpts, sigma): - return 0.5 * scipy.special.erfc((mo_energy_kpts - m) / sigma) - - def partition_occ(mo_occ, mo_energy_kpts): - mo_occ_kpts = [] - p1 = 0 - for e in mo_energy_kpts: - p0, p1 = p1, p1 + e.size - occ = mo_occ[p0:p1] - mo_occ_kpts.append(occ) - return mo_occ_kpts - - def get_occ(mo_energy_kpts=None, mo_coeff_kpts=None): +def _partition_occ(mo_occ, mo_energy_kpts): + mo_occ_kpts = [] + p1 = 0 + for e in mo_energy_kpts: + p0, p1 = p1, p1 + e.size + occ = mo_occ[p0:p1] + mo_occ_kpts.append(occ) + return mo_occ_kpts + +def _get_grad_tril(mo_coeff_kpts, mo_occ_kpts, fock): + grad_kpts = [] + for k, mo in enumerate(mo_coeff_kpts): + f_mo = reduce(numpy.dot, (mo.T.conj(), fock[k], mo)) + nmo = f_mo.shape[0] + grad_kpts.append(f_mo[numpy.tril_indices(nmo, -1)]) + return numpy.hstack(grad_kpts) + +class _SmearingKSCF(mol_addons._SmearingSCF): + def get_occ(self, mo_energy_kpts=None, mo_coeff_kpts=None): '''Label the occupancies for each orbital for sampled k-points. This is a k-point version of scf.hf.SCF.get_occ ''' - if is_rohf and fix_spin: - mo_energy_kpts=(mo_energy_kpts,mo_energy_kpts) - kpts = getattr(mf, 'kpts', None) - if isinstance(kpts, KPoints): - mo_energy_kpts = kpts.transform_mo_energy(mo_energy_kpts) - - #mo_occ_kpts = mf_class.get_occ(mf, mo_energy_kpts, mo_coeff_kpts) - if (mf.sigma == 0) or (not mf.sigma) or (not mf.smearing_method): - mo_occ_kpts = mf_class.get_occ(mf, mo_energy_kpts, mo_coeff_kpts) + from pyscf.scf import uhf, rohf, ghf + if (self.sigma == 0) or (not self.sigma) or (not self.smearing_method): + mo_occ_kpts = super().get_occ(mo_energy_kpts, mo_coeff_kpts) return mo_occ_kpts - if is_khf: - if isinstance(kpts, KPoints): - nkpts = kpts.nkpts - else: - nkpts = len(kpts) - else: - nkpts = 1 - if isinstance(mf.mol, pbcgto.Cell): - nelectron = mf.mol.tot_electrons(nkpts) - else: - nelectron = mf.mol.tot_electrons() - if is_uhf or (is_rohf and fix_spin): - nocc = nelectron - if fix_spin: - nocc = mf.nelec - mo_es = [] - mo_es.append(numpy.hstack(mo_energy_kpts[0])) - mo_es.append(numpy.hstack(mo_energy_kpts[1])) - else: - mo_es = numpy.append(numpy.hstack(mo_energy_kpts[0]), - numpy.hstack(mo_energy_kpts[1])) - elif is_ghf: - nocc = nelectron - mo_es = numpy.hstack(mo_energy_kpts) - else: - nocc = nelectron // 2 - mo_es = numpy.hstack(mo_energy_kpts) - - if mf.smearing_method.lower() == 'fermi': # Fermi-Dirac smearing - f_occ = fermi_smearing_occ - else: # Gaussian smearing - f_occ = gaussian_smearing_occ - - if fix_spin: - mo_energy = [] - mo_energy.append(numpy.sort(mo_es[0].ravel())) - mo_energy.append(numpy.sort(mo_es[1].ravel())) - else: - mo_energy = numpy.sort(mo_es.ravel()) + is_uhf = isinstance(self, uhf.UHF) + is_ghf = isinstance(self, ghf.GHF) + is_rhf = (not is_uhf) and (not is_ghf) + is_rohf = isinstance(self, rohf.ROHF) - sigma = mf.sigma - if fix_spin: - fermi = [mo_energy[0][nocc[0]-1], mo_energy[1][nocc[1]-1]] - else: - fermi = mo_energy[nocc-1] - if mu0 is None: - def nelec_cost_fn(m, _mo_es, _nelectron): - mo_occ_kpts = f_occ(m, _mo_es, sigma) - if is_rhf and not is_rohf: - mo_occ_kpts *= 2 - return (mo_occ_kpts.sum() - _nelectron)**2 - if fix_spin: - mu = [] - mo_occs = [] - res = scipy.optimize.minimize(nelec_cost_fn, fermi[0], args=(mo_es[0], nocc[0]), method='Powell', - options={'xtol': 1e-5, 'ftol': 1e-5, 'maxiter': 10000}) - mu.append(res.x) - mo_occs.append(f_occ(mu[0], mo_es[0], sigma)) - res = scipy.optimize.minimize(nelec_cost_fn, fermi[1], args=(mo_es[1], nocc[1]), method='Powell', - options={'xtol': 1e-5, 'ftol': 1e-5, 'maxiter': 10000}) - mu.append(res.x) - mo_occs.append(f_occ(mu[1], mo_es[1], sigma)) - f = copy.copy(mo_occs) - else: - res = scipy.optimize.minimize(nelec_cost_fn, fermi, args=(mo_es, nelectron), method='Powell', - options={'xtol': 1e-5, 'ftol': 1e-5, 'maxiter': 10000}) - mu = res.x - mo_occs = f = f_occ(mu, mo_es, sigma) + sigma = self.sigma + if self.smearing_method.lower() == 'fermi': + f_occ = mol_addons._fermi_smearing_occ else: - # If mu0 is given, fix mu instead of electron number. XXX -Chong Sun - mu = mu0 - mo_occs = f = f_occ(mu, mo_es, sigma) - - # See https://www.vasp.at/vasp-workshop/slides/k-points.pdf - if mf.smearing_method.lower() == 'fermi': - if fix_spin: - f[0] = f[0][(f[0]>0) & (f[0]<1)] - mf.entropy = -(f[0]*numpy.log(f[0]) + (1-f[0])*numpy.log(1-f[0])).sum() / nkpts - f[1] = f[1][(f[1]>0) & (f[1]<1)] - mf.entropy += -(f[1]*numpy.log(f[1]) + (1-f[1])*numpy.log(1-f[1])).sum() / nkpts - else: - f = f[(f>0) & (f<1)] - mf.entropy = -(f*numpy.log(f) + (1-f)*numpy.log(1-f)).sum() / nkpts + f_occ = mol_addons._gaussian_smearing_occ + + kpts = getattr(self, 'kpts', None) + if isinstance(kpts, KPoints): + nkpts = kpts.nkpts + mo_energy_kpts = kpts.transform_mo_energy(mo_energy_kpts) else: - if fix_spin: - mf.entropy = (numpy.exp(-((mo_es[0]-mu[0])/mf.sigma)**2).sum() - / (2*numpy.sqrt(numpy.pi)) / nkpts) - mf.entropy += (numpy.exp(-((mo_es[1]-mu[1])/mf.sigma)**2).sum() - / (2*numpy.sqrt(numpy.pi)) / nkpts) - else: - mf.entropy = (numpy.exp(-((mo_es-mu)/mf.sigma)**2).sum() - / (2*numpy.sqrt(numpy.pi)) / nkpts) - if is_rhf: - mo_occs *= 2 - mf.entropy *= 2 - - # DO NOT use numpy.array for mo_occ_kpts and mo_energy_kpts, they may - # have different dimensions for different k-points - if is_uhf: - if is_khf: - if fix_spin: - mo_occ_kpts =(partition_occ(mo_occs[0], mo_energy_kpts[0]), - partition_occ(mo_occs[1], mo_energy_kpts[1])) - else: - nao_tot = mo_occs.size//2 - mo_occ_kpts =(partition_occ(mo_occs[:nao_tot], mo_energy_kpts[0]), - partition_occ(mo_occs[nao_tot:], mo_energy_kpts[1])) - else: - mo_occ_kpts = partition_occ(mo_occs, mo_energy_kpts) - else: # rhf and ghf - if is_khf: - mo_occ_kpts = partition_occ(mo_occs, mo_energy_kpts) + nkpts = len(kpts) + + if self.fix_spin and (is_uhf or is_rohf): # spin separated fermi level + if is_rohf: # treat rohf as uhf + mo_energy_kpts = (mo_energy_kpts, mo_energy_kpts) + mo_es = [numpy.hstack(mo_energy_kpts[0]), + numpy.hstack(mo_energy_kpts[1])] + nocc = self.nelec + if self.mu0 is None: + mu_a, occa = mol_addons._smearing_optimize(f_occ, mo_es[0], nocc[0], sigma) + mu_b, occb = mol_addons._smearing_optimize(f_occ, mo_es[1], nocc[1], sigma) + mu = [mu_a, mu_b] + mo_occs = [occa, occb] else: - mo_occ_kpts = mo_occs - - if fix_spin: - logger.debug(mf, ' Alpha-spin Fermi level %g Sum mo_occ_kpts = %s should equal nelec = %s', + mu = self.mu0 + mo_occs = f_occ(mu[0], mo_es[0], sigma) + mo_occs = f_occ(mu[1], mo_es[1], sigma) + self.entropy = self._get_entropy(mo_es[0], mo_occs[0], mu[0]) + self.entropy += self._get_entropy(mo_es[1], mo_occs[1], mu[1]) + + fermi = (mol_addons._get_fermi_level(mo_es[0], nocc[0]), + mol_addons._get_fermi_level(mo_es[1], nocc[1])) + logger.debug(self, ' Alpha-spin Fermi level %g Sum mo_occ_kpts = %s should equal nelec = %s', fermi[0], mo_occs[0].sum(), nocc[0]) - logger.debug(mf, ' Beta-spin Fermi level %g Sum mo_occ_kpts = %s should equal nelec = %s', + logger.debug(self, ' Beta-spin Fermi level %g Sum mo_occ_kpts = %s should equal nelec = %s', fermi[1], mo_occs[1].sum(), nocc[1]) - logger.info(mf, ' sigma = %g Optimized mu_alpha = %.12g entropy = %.12g', - mf.sigma, mu[0], mf.entropy) - logger.info(mf, ' sigma = %g Optimized mu_beta = %.12g entropy = %.12g', - mf.sigma, mu[1], mf.entropy) + logger.info(self, ' sigma = %g Optimized mu_alpha = %.12g entropy = %.12g', + sigma, mu[0], self.entropy) + logger.info(self, ' sigma = %g Optimized mu_beta = %.12g entropy = %.12g', + sigma, mu[1], self.entropy) + + mo_occ_kpts =(_partition_occ(mo_occs[0], mo_energy_kpts[0]), + _partition_occ(mo_occs[1], mo_energy_kpts[1])) + tools.print_mo_energy_occ_kpts(self, mo_energy_kpts, mo_occ_kpts, True) + if is_rohf: + mo_occ_kpts = np.array(mo_occ_kpts, dtype=numpy.float64).sum(axis=0) else: - logger.debug(mf, ' Fermi level %g Sum mo_occ_kpts = %s should equal nelec = %s', + nocc = nelectron = self.mol.tot_electrons(nkpts) + if is_uhf: + mo_es_a = numpy.hstack(mo_energy_kpts[0]) + mo_es_b = numpy.hstack(mo_energy_kpts[1]) + mo_es = numpy.append(mo_es_a, mo_es_b) + else: + mo_es = numpy.hstack(mo_energy_kpts) + if is_rhf: + nocc = (nelectron + 1) // 2 + + if self.mu0 is None: + mu, mo_occs = mol_addons._smearing_optimize(f_occ, mo_es, nocc, sigma) + else: + # If mu0 is given, fix mu instead of electron number. XXX -Chong Sun + mu = self.mu0 + mo_occs = f_occ(mu, mo_es, sigma) + self.entropy = self._get_entropy(mo_es, mo_occs, mu) + if is_rhf: + mo_occs *= 2 + self.entropy *= 2 + + fermi = mol_addons._get_fermi_level(mo_es, nocc) + logger.debug(self, ' Fermi level %g Sum mo_occ_kpts = %s should equal nelec = %s', fermi, mo_occs.sum(), nelectron) - logger.info(mf, ' sigma = %g Optimized mu = %.12g entropy = %.12g', - mf.sigma, mu, mf.entropy) + logger.info(self, ' sigma = %g Optimized mu = %.12g entropy = %.12g', + sigma, mu, self.entropy) + + if is_uhf: + # mo_es_a and mo_es_b may have different dimensions for + # different k-points + nmo_a = mo_es_a.size + mo_occ_kpts =(_partition_occ(mo_occs[:nmo_a], mo_energy_kpts[0]), + _partition_occ(mo_occs[nmo_a:], mo_energy_kpts[1])) + else: + mo_occ_kpts = _partition_occ(mo_occs, mo_energy_kpts) + tools.print_mo_energy_occ_kpts(self, mo_energy_kpts, mo_occ_kpts, is_uhf) - kpts = getattr(mf, 'kpts', None) if isinstance(kpts, KPoints): if is_uhf: mo_occ_kpts = (kpts.check_mo_occ_symmetry(mo_occ_kpts[0]), kpts.check_mo_occ_symmetry(mo_occ_kpts[1])) else: mo_occ_kpts = kpts.check_mo_occ_symmetry(mo_occ_kpts) - - if is_khf: - tools.print_mo_energy_occ_kpts(mf,mo_energy_kpts,mo_occ_kpts,is_uhf) - else: - tools.print_mo_energy_occ(mf,mo_energy_kpts,mo_occ_kpts,is_uhf) - if is_rohf and fix_spin: - mo_occ_kpts=mo_occ_kpts[0]+mo_occ_kpts[1] return mo_occ_kpts - def get_grad_tril(mo_coeff_kpts, mo_occ_kpts, fock): - if is_khf: - grad_kpts = [] - for k, mo in enumerate(mo_coeff_kpts): - f_mo = reduce(numpy.dot, (mo.T.conj(), fock[k], mo)) - nmo = f_mo.shape[0] - grad_kpts.append(f_mo[numpy.tril_indices(nmo, -1)]) - return numpy.hstack(grad_kpts) - else: - f_mo = reduce(numpy.dot, (mo_coeff_kpts.T.conj(), fock, mo_coeff_kpts)) - nmo = f_mo.shape[0] - return f_mo[numpy.tril_indices(nmo, -1)] + def get_grad(self, mo_coeff_kpts, mo_occ_kpts, fock=None): + from pyscf.scf import uhf + if (self.sigma == 0) or (not self.sigma) or (not self.smearing_method): + return super().get_grad(mo_coeff_kpts, mo_occ_kpts, fock) - def get_grad(mo_coeff_kpts, mo_occ_kpts, fock=None): - if (mf.sigma == 0) or (not mf.sigma) or (not mf.smearing_method): - return mf_class.get_grad(mf, mo_coeff_kpts, mo_occ_kpts, fock) if fock is None: - dm1 = mf.make_rdm1(mo_coeff_kpts, mo_occ_kpts) - fock = mf.get_hcore() + mf.get_veff(mf.mol, dm1) - if is_uhf: - ga = get_grad_tril(mo_coeff_kpts[0], mo_occ_kpts[0], fock[0]) - gb = get_grad_tril(mo_coeff_kpts[1], mo_occ_kpts[1], fock[1]) + dm1 = self.make_rdm1(mo_coeff_kpts, mo_occ_kpts) + fock = self.get_hcore() + self.get_veff(self.mol, dm1) + if isinstance(self, uhf.UHF): + ga = _get_grad_tril(mo_coeff_kpts[0], mo_occ_kpts[0], fock[0]) + gb = _get_grad_tril(mo_coeff_kpts[1], mo_occ_kpts[1], fock[1]) return numpy.hstack((ga,gb)) else: # rhf and ghf - return get_grad_tril(mo_coeff_kpts, mo_occ_kpts, fock) - - def energy_tot(dm=None, h1e=None, vhf=None): - e_tot = mf.energy_elec(dm, h1e, vhf)[0] + mf.energy_nuc() - if (mf.sigma and mf.smearing_method and - mf.entropy is not None): - mf.e_free = e_tot - mf.sigma * mf.entropy - mf.e_zero = e_tot - mf.sigma * mf.entropy * .5 - logger.info(mf, ' Total E(T) = %.15g Free energy = %.15g E0 = %.15g', - e_tot, mf.e_free, mf.e_zero) - return e_tot - - mf.sigma = sigma - mf.smearing_method = method - mf.entropy = None - mf.e_free = None - mf.e_zero = None - mf._keys = set(['sigma', 'smearing_method', - 'entropy', 'e_free', 'e_zero']) + return _get_grad_tril(mo_coeff_kpts, mo_occ_kpts, fock) - mf.get_occ = get_occ - mf.energy_tot = energy_tot - mf.get_grad = get_grad +def smearing(mf, sigma=None, method=SMEARING_METHOD, mu0=None, fix_spin=False): + '''Fermi-Dirac or Gaussian smearing''' + from pyscf.pbc.scf import khf + if not isinstance(mf, khf.KSCF): + return mol_addons.smearing(mf, sigma, method, mu0, fix_spin) + + if isinstance(mf, mol_addons._SmearingSCF): + mf.sigma = sigma + mf.smearing_method = method + mf.mu0 = mu0 + mf.fix_spin = fix_spin + return mf + + return lib.set_class(_SmearingKSCF(mf, sigma, method, mu0, fix_spin), + (_SmearingKSCF, mf.__class__)) + +def smearing_(mf, *args, **kwargs): + mf1 = smearing(mf, *args, **kwargs) + mf.__class__ = mf1.__class__ + mf.__dict__.update(mf1.__dict__) return mf def canonical_occ_(mf, nelec=None): @@ -617,31 +533,3 @@ def _get_moek(kC, kfao): else: log.error(f'Unknown SCF type {mf.__class__.__name__}') raise NotImplementedError - - -del (SMEARING_METHOD) - - -if __name__ == '__main__': - import pyscf.pbc.scf as pscf - cell = pbcgto.Cell() - cell.atom = ''' - He 0 0 1 - He 1 0 1 - ''' - cell.basis = 'ccpvdz' - cell.a = numpy.eye(3) * 4 - cell.mesh = [17] * 3 - cell.verbose = 3 - cell.build() - nks = [2,1,1] - mf = pscf.KUHF(cell, cell.make_kpts(nks)).density_fit() - mf = smearing_(mf, .1) - mf.kernel() - print(mf.e_tot - -5.56769351866668) - mf = smearing_(mf, .1, mu0=0.351195741757) - mf.kernel() - print(mf.e_tot - -5.56769351866668) - mf = smearing_(mf, .1, method='gauss') - mf.kernel() - print(mf.e_tot - -5.56785857886738) diff --git a/pyscf/pbc/scf/test/test_addons.py b/pyscf/pbc/scf/test/test_addons.py index 197552e9fc..dd9e9a0964 100644 --- a/pyscf/pbc/scf/test/test_addons.py +++ b/pyscf/pbc/scf/test/test_addons.py @@ -81,6 +81,29 @@ def test_kuhf_smearing(self): occ = mf.get_occ(mo_energy_kpts) self.assertAlmostEqual(mf.entropy, 0.9554526863670467/2, 9) + def test_kuhf_smearing1(self): + cell = pbcgto.Cell() + cell.atom = ''' + He 0 0 1 + He 1 0 1 + ''' + cell.basis = 'ccpvdz' + cell.a = numpy.eye(3) * 4 + cell.precision = 1e-6 + cell.verbose = 3 + cell.build() + nks = [2,1,1] + mf = pscf.KUHF(cell, cell.make_kpts(nks)).density_fit() + mf = smearing_(mf, .1) + mf.kernel() + self.assertAlmostEqual(mf.e_tot, -5.56769351866668, 6) + mf = smearing_(mf, .1, mu0=0.351195741757) + mf.kernel() + self.assertAlmostEqual(mf.e_tot, -5.56769351866668, 6) + mf = smearing_(mf, .1, method='gauss') + mf.kernel() + self.assertAlmostEqual(mf.e_tot, -5.56785857886738, 6) + def test_rhf_smearing(self): mf = pscf.RHF(cell) pscf.addons.smearing_(mf, 0.1, 'fermi') @@ -376,27 +399,6 @@ def get_veff(cell, dm, *args): self.assertAlmostEqual(kmf.e_free, 1.3942942592412, 9) self.assertAlmostEqual(lib.fp(kmf.mo_occ), 0.035214493032250, 9) - def test_rohf_smearing(self): - from pyscf import gto, scf - # Fe2 from https://doi.org/10.1021/acs.jpca.1c05769 - mol = gto.M( - atom=''' - Fe 0. 0. 0. - Fe 2.01 0. 0. - ''', basis="lanl2dz", ecp="lanl2dz", symmetry=False, unit='Angstrom', - spin=6, charge=0, verbose=0) - myhf_s = scf.ROHF(mol) - myhf_s = pscf.addons.smearing_(myhf_s, sigma=0.01, method='gaussian', fix_spin=True) - myhf_s.kernel() - # macos py3.7 CI -242.4828982467762 - # linux py3.11 CI -242.48289824670388 - self.assertAlmostEqual(myhf_s.e_tot, -242.482898246, 6) - self.assertAlmostEqual(myhf_s.entropy, 0.45197, 4) - myhf2 = scf.ROHF(mol).newton() - myhf2.kernel(myhf_s.make_rdm1()) - # note 16mHa lower energy than myhf - self.assertAlmostEqual(myhf2.e_tot, -244.9808750, 6) - if __name__ == '__main__': print("Full Tests for pbc.scf.addons") diff --git a/pyscf/scf/addons.py b/pyscf/scf/addons.py index dcbcf9ee6c..ba3ecddd63 100644 --- a/pyscf/scf/addons.py +++ b/pyscf/scf/addons.py @@ -34,9 +34,203 @@ FORCE_PIVOTED_CHOLESKY = getattr(__config__, 'scf_addons_force_cholesky', False) LINEAR_DEP_TRIGGER = getattr(__config__, 'scf_addons_remove_linear_dep_trigger', 1e-10) -def smearing_(*args, **kwargs): - from pyscf.pbc.scf.addons import smearing_ - return smearing_(*args, **kwargs) +SMEARING_METHOD = getattr(__config__, 'pbc_scf_addons_smearing_method', 'fermi') + +def smearing(mf, sigma=None, method=SMEARING_METHOD, mu0=None, fix_spin=False): + '''Fermi-Dirac or Gaussian smearing''' + if isinstance(mf, _SmearingSCF): + mf.sigma = sigma + mf.method = method + mf.mu0 = mu0 + mf.fix_spin = fix_spin + return mf + + return lib.set_class(_SmearingSCF(mf, sigma, method, mu0, fix_spin), + (_SmearingSCF, mf.__class__)) + +def smearing_(mf, *args, **kwargs): + mf1 = smearing(mf, *args, **kwargs) + mf.__class__ = mf1.__class__ + mf.__dict__.update(mf1.__dict__) + return mf + +def _get_grad_tril(mo_coeff, mo_occ, fock): + f_mo = reduce(numpy.dot, (mo_coeff.T.conj(), fock, mo_coeff)) + return f_mo[numpy.tril_indices_from(f_mo, -1)] + +def _fermi_smearing_occ(mu, mo_energy, sigma): + '''Fermi-Dirac smearing''' + occ = numpy.zeros_like(mo_energy) + de = (mo_energy - mu) / sigma + occ[de<40] = 1./(numpy.exp(de[de<40])+1.) + return occ + +def _gaussian_smearing_occ(mu, mo_energy, sigma): + '''Gaussian smearing''' + return 0.5 * scipy.special.erfc((mo_energy - mu) / sigma) + +def _smearing_optimize(f_occ, mo_es, nocc, sigma): + def nelec_cost_fn(m, mo_es, nocc): + mo_occ = f_occ(m, mo_es, sigma) + return (mo_occ.sum() - nocc)**2 + + fermi = _get_fermi_level(mo_es, nocc) + res = scipy.optimize.minimize( + nelec_cost_fn, fermi, args=(mo_es, nocc), method='Powell', + options={'xtol': 1e-5, 'ftol': 1e-5, 'maxiter': 10000}) + mu = res.x + mo_occs = f_occ(mu, mo_es, sigma) + return mu, mo_occs + +def _get_fermi_level(mo_energy, nocc): + mo_e_sorted = numpy.sort(mo_energy) + return mo_e_sorted[nocc-1] + +class _SmearingSCF: + + __name_mixin__ = 'Smearing' + + _keys = set([ + 'sigma', 'smearing_method', 'mu0', 'fix_spin', 'entropy', 'e_free', 'e_zero' + ]) + + def __init__(self, mf, sigma, method, mu0, fix_spin): + self.__dict__.update(mf.__dict__) + self.sigma = sigma + self.smearing_method = method + self.mu0 = mu0 + self.fix_spin = fix_spin + self.entropy = None + self.e_free = None + self.e_zero = None + + def undo_smearing(self): + obj = lib.view(self, lib.drop_class(self.__class__, _SmearingSCF)) + del obj.sigma + del obj.smearing_method + del obj.fix_spin + del obj.entropy + del obj.e_free + del obj.e_zero + return obj + + def get_occ(self, mo_energy=None, mo_coeff=None): + '''Label the occupancies for each orbital + ''' + from pyscf.scf import uhf, rohf, ghf + from pyscf.pbc.tools import print_mo_energy_occ + if (self.sigma == 0) or (not self.sigma) or (not self.smearing_method): + mo_occ = super().get_occ(mo_energy, mo_coeff) + return mo_occ + + is_uhf = isinstance(self, uhf.UHF) + is_ghf = isinstance(self, ghf.GHF) + is_rhf = (not is_uhf) and (not is_ghf) + is_rohf = isinstance(self, rohf.ROHF) + + sigma = self.sigma + if self.smearing_method.lower() == 'fermi': + f_occ = _fermi_smearing_occ + else: + f_occ = _gaussian_smearing_occ + + if self.fix_spin and (is_uhf or is_rohf): # spin separated fermi level + if is_rohf: # treat rohf as uhf + mo_es = (mo_energy, mo_energy) + else: + mo_es = mo_energy + nocc = self.nelec + if self.mu0 is None: + mu_a, occa = _smearing_optimize(f_occ, mo_es[0], nocc[0], sigma) + mu_b, occb = _smearing_optimize(f_occ, mo_es[1], nocc[1], sigma) + mu = [mu_a, mu_b] + mo_occs = [occa, occb] + else: + mu = self.mu0 + mo_occs = f_occ(mu[0], mo_es[0], sigma) + mo_occs = f_occ(mu[1], mo_es[1], sigma) + self.entropy = self._get_entropy(mo_es[0], mo_occs[0], mu[0]) + self.entropy += self._get_entropy(mo_es[1], mo_occs[1], mu[1]) + fermi = (_get_fermi_level(mo_es[0], nocc[0]), + _get_fermi_level(mo_es[1], nocc[1])) + + logger.debug(self, ' Alpha-spin Fermi level %g Sum mo_occ = %s should equal nelec = %s', + fermi[0], mo_occs[0].sum(), nocc[0]) + logger.debug(self, ' Beta-spin Fermi level %g Sum mo_occ = %s should equal nelec = %s', + fermi[1], mo_occs[1].sum(), nocc[1]) + logger.info(self, ' sigma = %g Optimized mu_alpha = %.12g entropy = %.12g', + sigma, mu[0], self.entropy) + logger.info(self, ' sigma = %g Optimized mu_beta = %.12g entropy = %.12g', + sigma, mu[1], self.entropy) + if self.verbose >= logger.DEBUG: + print_mo_energy_occ(self, mo_energy, mo_occs, True) + if is_rohf: + mo_occs = mo_occs[0] + mo_occs[1] + else: # all orbitals treated with the same fermi level + nocc = nelectron = self.mol.tot_electrons() + if is_uhf: + mo_es = numpy.hstack(mo_energy) + else: + mo_es = mo_energy + if is_rhf: + nocc = (nelectron + 1) // 2 + + if self.mu0 is None: + mu, mo_occs = _smearing_optimize(f_occ, mo_es, nocc, sigma) + else: + # If mu0 is given, fix mu instead of electron number. XXX -Chong Sun + mu = self.mu0 + mo_occs = f_occ(mu, mo_es, sigma) + self.entropy = self._get_entropy(mo_es, mo_occs, mu) + if is_rhf: + mo_occs *= 2 + self.entropy *= 2 + + fermi = _get_fermi_level(mo_es, nocc) + logger.debug(self, ' Fermi level %g Sum mo_occ = %s should equal nelec = %s', + fermi, mo_occs.sum(), nelectron) + logger.info(self, ' sigma = %g Optimized mu = %.12g entropy = %.12g', + sigma, mu, self.entropy) + if is_uhf: + mo_occs = mo_occs.reshape(2, -1) + if self.verbose >= logger.DEBUG: + print_mo_energy_occ(self, mo_energy, mo_occs, is_uhf) + return mo_occs + + # See https://www.vasp.at/vasp-workshop/slides/k-points.pdf + def _get_entropy(self, mo_energy, mo_occ, mu): + if self.smearing_method.lower() == 'fermi': + f = mo_occ + f = f[(f>0) & (f<1)] + entropy = -(f*numpy.log(f) + (1-f)*numpy.log(1-f)).sum() + else: + entropy = (numpy.exp(-((mo_energy-mu)/self.sigma)**2).sum() + / (2*numpy.sqrt(numpy.pi))) + return entropy + + def get_grad(self, mo_coeff, mo_occ, fock=None): + from pyscf.scf import uhf + if (self.sigma == 0) or (not self.sigma) or (not self.smearing_method): + return super().get_grad(mo_coeff, mo_occ, fock) + + if fock is None: + dm1 = self.make_rdm1(mo_coeff, mo_occ) + fock = self.get_hcore() + self.get_veff(self.mol, dm1) + if isinstance(self, uhf.UHF): + ga = _get_grad_tril(mo_coeff[0], mo_occ[0], fock[0]) + gb = _get_grad_tril(mo_coeff[1], mo_occ[1], fock[1]) + return numpy.hstack((ga,gb)) + else: # rhf and ghf + return _get_grad_tril(mo_coeff, mo_occ, fock) + + def energy_tot(self, dm=None, h1e=None, vhf=None): + e_tot = self.energy_elec(dm, h1e, vhf)[0] + self.energy_nuc() + if self.sigma and self.smearing_method and self.entropy is not None: + self.e_free = e_tot - self.sigma * self.entropy + self.e_zero = e_tot - self.sigma * self.entropy * .5 + logger.info(self, ' Total E(T) = %.15g Free energy = %.15g E0 = %.15g', + e_tot, self.e_free, self.e_zero) + return e_tot def frac_occ_(mf, tol=1e-3): ''' diff --git a/pyscf/scf/test/test_addons.py b/pyscf/scf/test/test_addons.py index 8249062e14..3c5eaed2a1 100644 --- a/pyscf/scf/test/test_addons.py +++ b/pyscf/scf/test/test_addons.py @@ -414,6 +414,51 @@ def test_remove_lindep(self): lindep=1e-9).run() self.assertAlmostEqual(mf.e_tot, -1.6291001503057689, 7) + def test_rohf_smearing(self): + # Fe2 from https://doi.org/10.1021/acs.jpca.1c05769 + mol = gto.M( + atom=''' + Fe 0. 0. 0. + Fe 2.01 0. 0. + ''', basis="lanl2dz", ecp="lanl2dz", symmetry=False, unit='Angstrom', + spin=6, charge=0, verbose=0) + myhf_s = scf.ROHF(mol) + myhf_s = addons.smearing(myhf_s, sigma=0.01, method='gaussian', fix_spin=True) + myhf_s.kernel() + # macos py3.7 CI -242.4828982467762 + # linux py3.11 CI -242.48289824670388 + self.assertAlmostEqual(myhf_s.e_tot, -242.482898246, 6) + self.assertAlmostEqual(myhf_s.entropy, 0.45197, 4) + myhf2 = myhf_s.undo_smearing().newton() + myhf2.kernel(myhf_s.make_rdm1()) + # FIXME: note 16mHa lower energy than myhf + self.assertAlmostEqual(myhf2.e_tot, -244.9808750, 6) + + myhf_s.smearing_method = 'fermi' + myhf_s.kernel() + self.assertAlmostEqual(myhf_s.e_tot, -244.200255453, 6) + self.assertAlmostEqual(myhf_s.entropy, 3.585155, 4) + + myhf_s = scf.UHF(mol) + myhf_s = addons.smearing_(myhf_s, sigma=0.01, method='fermi', fix_spin=True) + myhf_s.kernel() + self.assertAlmostEqual(myhf_s.e_tot, -244.9873314, 5) + self.assertAlmostEqual(myhf_s.entropy, 0, 3) + + myhf_s = scf.UHF(mol) + myhf_s = addons.smearing_(myhf_s, sigma=0.01, method='fermi', fix_spin=True) + myhf_s.sigma = 0.1 + myhf_s.fix_spin = False + myhf_s.kernel() + self.assertAlmostEqual(myhf_s.e_tot, -243.086989253, 6) + self.assertAlmostEqual(myhf_s.entropy, 17.11431, 4) + self.assertTrue(myhf_s.converged) + + myhf_s.mu = -0.2482816 + myhf_s.kernel(dm0=myhf_s.make_rdm1()) + self.assertAlmostEqual(myhf_s.e_tot, -243.086989253, 6) + self.assertAlmostEqual(myhf_s.entropy, 17.11431, 4) + if __name__ == "__main__": print("Full Tests for addons") unittest.main()