|
| 1 | +# Copyright 2021-2025 The PySCF Developers. All Rights Reserved. |
| 2 | +# |
| 3 | +# Licensed under the Apache License, Version 2.0 (the "License"); |
| 4 | +# you may not use this file except in compliance with the License. |
| 5 | +# You may obtain a copy of the License at |
| 6 | +# |
| 7 | +# http://www.apache.org/licenses/LICENSE-2.0 |
| 8 | +# |
| 9 | +# Unless required by applicable law or agreed to in writing, software |
| 10 | +# distributed under the License is distributed on an "AS IS" BASIS, |
| 11 | +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. |
| 12 | +# See the License for the specific language governing permissions and |
| 13 | +# limitations under the License. |
| 14 | + |
| 15 | +from functools import reduce |
| 16 | +import numpy as np |
| 17 | +import cupy as cp |
| 18 | +from pyscf import lib |
| 19 | +from pyscf.pbc.lib.kpts import KPoints |
| 20 | +from gpu4pyscf.lib import logger |
| 21 | +from gpu4pyscf.scf import smearing as mol_smearing |
| 22 | +from gpu4pyscf.lib.cupy_helper import contract |
| 23 | + |
| 24 | +SMEARING_METHOD = mol_smearing.SMEARING_METHOD |
| 25 | + |
| 26 | +def smearing(mf, sigma=None, method=None, mu0=None, fix_spin=False): |
| 27 | + '''Fermi-Dirac or Gaussian smearing''' |
| 28 | + from gpu4pyscf.pbc.scf import khf |
| 29 | + if not isinstance(mf, khf.KSCF): |
| 30 | + return mol_smearing.smearing(mf, sigma, method, mu0, fix_spin) |
| 31 | + |
| 32 | + if isinstance(mf, mol_smearing._SmearingSCF): |
| 33 | + mf.sigma = sigma |
| 34 | + mf.smearing_method = method |
| 35 | + mf.mu0 = mu0 |
| 36 | + mf.fix_spin = fix_spin |
| 37 | + return mf |
| 38 | + |
| 39 | + return lib.set_class(_SmearingKSCF(mf, sigma, method, mu0, fix_spin), |
| 40 | + (_SmearingKSCF, mf.__class__)) |
| 41 | + |
| 42 | +def _partition_occ(mo_occ, mo_energy_kpts): |
| 43 | + dims = [e.size for e in mo_energy_kpts] |
| 44 | + offsets = np.cumsum(dims) |
| 45 | + mo_occ_kpts = np.split(mo_occ, offsets[:-1]) |
| 46 | + return mo_occ_kpts |
| 47 | + |
| 48 | +def _get_grad_tril(mo_coeff_kpts, mo_occ_kpts, fock): |
| 49 | + nmo = mo_occ_kpts.shape[-1] |
| 50 | + i, j = cp.tril_indices(nmo, -1) |
| 51 | + fc = contract('kpq,kqj->kpj', fock, mo_coeff_kpts) |
| 52 | + grad_kpts = contract('kpi,kpj->kij', mo_coeff_kpts.conj(), fc) |
| 53 | + return grad_kpts[:,i,j].ravel() |
| 54 | + |
| 55 | +class _SmearingKSCF(mol_smearing._SmearingSCF): |
| 56 | + def get_occ(self, mo_energy_kpts=None, mo_coeff_kpts=None): |
| 57 | + '''Label the occupancies for each orbital for sampled k-points. |
| 58 | +
|
| 59 | + This is a k-point version of scf.hf.SCF.get_occ |
| 60 | + ''' |
| 61 | + from gpu4pyscf.pbc import scf |
| 62 | + if (self.sigma == 0) or (not self.sigma) or (not self.smearing_method): |
| 63 | + mo_occ_kpts = super().get_occ(mo_energy_kpts, mo_coeff_kpts) |
| 64 | + return mo_occ_kpts |
| 65 | + |
| 66 | + is_uhf = self.istype('KUHF') |
| 67 | + is_rhf = self.istype('KRHF') |
| 68 | + |
| 69 | + sigma = self.sigma |
| 70 | + if self.smearing_method.lower() == 'fermi': |
| 71 | + f_occ = mol_smearing._fermi_smearing_occ |
| 72 | + else: |
| 73 | + f_occ = mol_smearing._gaussian_smearing_occ |
| 74 | + |
| 75 | + kpts = getattr(self, 'kpts', None) |
| 76 | + if isinstance(kpts, KPoints): |
| 77 | + raise NotImplementedError |
| 78 | + else: |
| 79 | + nkpts = len(kpts) |
| 80 | + |
| 81 | + mo_energy_kpts = mo_energy_kpts.get() |
| 82 | + if self.fix_spin and is_uhf: # spin separated fermi level |
| 83 | + mo_es = mo_energy_kpts.reshape(2, -1) |
| 84 | + nocc = self.nelec |
| 85 | + if self.mu0 is None: |
| 86 | + mu_a, occa = mol_smearing._smearing_optimize(f_occ, mo_es[0], nocc[0], sigma) |
| 87 | + mu_b, occb = mol_smearing._smearing_optimize(f_occ, mo_es[1], nocc[1], sigma) |
| 88 | + else: |
| 89 | + if np.isscalar(self.mu0): |
| 90 | + mu_a = mu_b = self.mu0 |
| 91 | + elif len(self.mu0) == 2: |
| 92 | + mu_a, mu_b = self.mu0 |
| 93 | + else: |
| 94 | + raise TypeError(f'Unsupported mu0: {self.mu0}') |
| 95 | + occa = f_occ(mu_a, mo_es[0], sigma) |
| 96 | + occb = f_occ(mu_b, mo_es[1], sigma) |
| 97 | + mu = [mu_a, mu_b] |
| 98 | + mo_occs = [occa, occb] |
| 99 | + self.entropy = self._get_entropy(mo_es[0], mo_occs[0], mu[0]) |
| 100 | + self.entropy += self._get_entropy(mo_es[1], mo_occs[1], mu[1]) |
| 101 | + self.entropy /= nkpts |
| 102 | + |
| 103 | + if self.verbose >= logger.INFO: |
| 104 | + fermi = (mol_smearing._get_fermi(mo_es[0], nocc[0]), |
| 105 | + mol_smearing._get_fermi(mo_es[1], nocc[1])) |
| 106 | + logger.debug(self, ' Alpha-spin Fermi level %g Sum mo_occ_kpts = %s should equal nelec = %s', |
| 107 | + fermi[0], mo_occs[0].sum(), nocc[0]) |
| 108 | + logger.debug(self, ' Beta-spin Fermi level %g Sum mo_occ_kpts = %s should equal nelec = %s', |
| 109 | + fermi[1], mo_occs[1].sum(), nocc[1]) |
| 110 | + logger.info(self, ' sigma = %g Optimized mu_alpha = %.12g entropy = %.12g', |
| 111 | + sigma, mu[0], self.entropy) |
| 112 | + logger.info(self, ' sigma = %g Optimized mu_beta = %.12g entropy = %.12g', |
| 113 | + sigma, mu[1], self.entropy) |
| 114 | + |
| 115 | + mo_occ_kpts =(_partition_occ(mo_occs[0], mo_energy_kpts[0]), |
| 116 | + _partition_occ(mo_occs[1], mo_energy_kpts[1])) |
| 117 | + else: |
| 118 | + nocc = nelectron = self.mol.tot_electrons(nkpts) |
| 119 | + mo_es = mo_energy_kpts.ravel() |
| 120 | + if is_rhf: |
| 121 | + nocc = (nelectron + 1) // 2 |
| 122 | + |
| 123 | + if self.mu0 is None: |
| 124 | + mu, mo_occs = mol_smearing._smearing_optimize(f_occ, mo_es, nocc, sigma) |
| 125 | + else: |
| 126 | + # If mu0 is given, fix mu instead of electron number. XXX -Chong Sun |
| 127 | + mu = self.mu0 |
| 128 | + assert np.isscalar(mu) |
| 129 | + mo_occs = f_occ(mu, mo_es, sigma) |
| 130 | + self.entropy = self._get_entropy(mo_es, mo_occs, mu) / nkpts |
| 131 | + if is_rhf: |
| 132 | + mo_occs *= 2 |
| 133 | + self.entropy *= 2 |
| 134 | + |
| 135 | + if self.verbose >= logger.INFO: |
| 136 | + fermi = mol_smearing._get_fermi(mo_es, nocc) |
| 137 | + logger.debug(self, ' Fermi level %g Sum mo_occ_kpts = %s should equal nelec = %s', |
| 138 | + fermi, mo_occs.sum(), nelectron) |
| 139 | + logger.info(self, ' sigma = %g Optimized mu = %.12g entropy = %.12g', |
| 140 | + sigma, mu, self.entropy) |
| 141 | + |
| 142 | + if is_uhf: |
| 143 | + # mo_es_a and mo_es_b may have different dimensions for |
| 144 | + # different k-points |
| 145 | + nmo_a = mo_energy_kpts[0].size |
| 146 | + mo_occ_kpts = (_partition_occ(mo_occs[:nmo_a], mo_energy_kpts[0]), |
| 147 | + _partition_occ(mo_occs[nmo_a:], mo_energy_kpts[1])) |
| 148 | + else: |
| 149 | + mo_occ_kpts = _partition_occ(mo_occs, mo_energy_kpts) |
| 150 | + return cp.array(mo_occ_kpts) |
| 151 | + |
| 152 | + def get_grad(self, mo_coeff_kpts, mo_occ_kpts, fock=None): |
| 153 | + if (self.sigma == 0) or (not self.sigma) or (not self.smearing_method): |
| 154 | + return super().get_grad(mo_coeff_kpts, mo_occ_kpts, fock) |
| 155 | + |
| 156 | + if fock is None: |
| 157 | + dm1 = self.make_rdm1(mo_coeff_kpts, mo_occ_kpts) |
| 158 | + fock = self.get_hcore() + self.get_veff(self.mol, dm1) |
| 159 | + if self.istype('KUHF'): |
| 160 | + ga = _get_grad_tril(mo_coeff_kpts[0], mo_occ_kpts[0], fock[0]) |
| 161 | + gb = _get_grad_tril(mo_coeff_kpts[1], mo_occ_kpts[1], fock[1]) |
| 162 | + return cp.hstack((ga, gb)) |
| 163 | + else: # rhf and ghf |
| 164 | + return _get_grad_tril(mo_coeff_kpts, mo_occ_kpts, fock) |
| 165 | + |
| 166 | + def to_cpu(self): |
| 167 | + from pyscf.pbc.scf.addons import smearing |
| 168 | + return smearing(self.undo_smearing().to_cpu(), self.sigma, |
| 169 | + self.smearing_method, self.mu0, self.fix_spin) |
| 170 | + |
| 171 | +def from_cpu(method): |
| 172 | + from pyscf.scf.addons import _SmearingSCF |
| 173 | + assert isinstance(method, _SmearingSCF) |
| 174 | + return smearing(method.undo_smearing().to_cpu(), method.sigma, |
| 175 | + method.smearing_method, method.mu0, method.fix_spin) |
0 commit comments