#!/usr/bin/env python
# Copyright 2020-2023 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: Xing Zhang <zhangxing.nju@gmail.com>
#
import numpy as np
import h5py
import scipy.linalg
from pyscf import __config__
from pyscf import lib
from pyscf.lib import logger
from pyscf.scf import hf as mol_hf
from pyscf.pbc import tools
from pyscf.pbc.lib import kpts as libkpts
from pyscf.pbc.lib.kpts_helper import is_trim
from pyscf.pbc.scf import khf
from pyscf.pbc.scf import hf as pbchf
from pyscf.symm import symmetrize_matrix
from pyscf.scf.hf_symm import so2ao_mo_coeff
[docs]
@lib.with_doc(khf.get_occ.__doc__)
def get_occ(mf, mo_energy_kpts=None, mo_coeff_kpts=None):
if mo_energy_kpts is None:
mo_energy_kpts = mf.mo_energy
cell = mf.cell
kpts = mf.kpts
assert isinstance(kpts, libkpts.KPoints)
nocc = cell.tot_electrons(kpts.nkpts) // 2
mo_energy_kpts = kpts.transform_mo_energy(mo_energy_kpts)
mo_energy = np.sort(np.hstack(mo_energy_kpts))
fermi = mo_energy[nocc-1]
mo_occ_kpts = []
for mo_e in mo_energy_kpts:
mo_occ_kpts.append((mo_e <= fermi).astype(np.double) * 2)
if nocc < mo_energy.size:
logger.info(mf, 'HOMO = %.12g LUMO = %.12g',
mo_energy[nocc-1], mo_energy[nocc])
if mo_energy[nocc-1]+1e-3 > mo_energy[nocc]:
logger.warn(mf, 'HOMO %.12g == LUMO %.12g',
mo_energy[nocc-1], mo_energy[nocc])
else:
logger.info(mf, 'HOMO = %.12g', mo_energy[nocc-1])
if mf.verbose >= logger.DEBUG:
np.set_printoptions(threshold=len(mo_energy))
logger.debug(mf, ' k-point mo_energy')
for k,kpt in enumerate(mf.cell.get_scaled_kpts(mf.kpts, kpts_in_ibz=False)):
logger.debug(mf, ' %2d (%6.3f %6.3f %6.3f) %s %s',
k, kpt[0], kpt[1], kpt[2],
np.sort(mo_energy_kpts[k][mo_occ_kpts[k]> 0]),
np.sort(mo_energy_kpts[k][mo_occ_kpts[k]==0]))
np.set_printoptions(threshold=1000)
mo_occ_kpts = kpts.check_mo_occ_symmetry(mo_occ_kpts)
return np.array(mo_occ_kpts)
[docs]
@lib.with_doc(khf.energy_elec.__doc__)
def energy_elec(mf, dm_kpts=None, h1e_kpts=None, vhf_kpts=None):
if dm_kpts is None: dm_kpts = mf.make_rdm1()
if h1e_kpts is None: h1e_kpts = mf.get_hcore()
if vhf_kpts is None: vhf_kpts = mf.get_veff(mf.cell, dm_kpts)
kpts_weights = mf.kpts.weights_ibz
e1 = np.einsum('k,kij,kji', kpts_weights, dm_kpts, h1e_kpts)
e_coul = np.einsum('k,kij,kji', kpts_weights, dm_kpts, vhf_kpts) * 0.5
mf.scf_summary['e1'] = e1.real
mf.scf_summary['e2'] = e_coul.real
logger.debug(mf, 'E1 = %s E_coul = %s', e1, e_coul)
if khf.CHECK_COULOMB_IMAG and abs(e_coul.imag) > mf.cell.precision*10:
logger.warn(mf, "Coulomb energy has imaginary part %s. "
"Coulomb integrals (e-e, e-N) may not converge !",
e_coul.imag)
return (e1+e_coul).real, e_coul.real
[docs]
@lib.with_doc(khf.get_rho.__doc__)
def get_rho(mf, dm=None, grids=None, kpts=None):
if isinstance(kpts, np.ndarray):
return khf.get_rho(mf, dm, grids, kpts)
if dm is None: dm = mf.make_rdm1()
if kpts is None: kpts = mf.kpts
if isinstance(dm[0], np.ndarray) and dm[0].ndim == 3:
ndm = len(dm[0])
else:
ndm = len(dm)
if ndm != kpts.nkpts_ibz:
raise RuntimeError("Number of input density matrices does not \
match the number of IBZ kpts: %d vs %d."
% (ndm, kpts.nkpts_ibz))
dm = kpts.transform_dm(dm)
return khf.get_rho(mf, dm, grids, kpts.kpts)
[docs]
def eig(kmf, h_kpts, s_kpts, overwrite=False, x=None):
from pyscf.scf.hf_symm import eig as eig_symm
cell = kmf.cell
symm_orb = cell.symm_orb
irrep_id = cell.irrep_id
nkpts, nao = h_kpts.shape[:2]
assert len(symm_orb) == nkpts
eig_kpts = []
mo_coeff_kpts = []
for k in range(nkpts):
e, c = eig_symm(kmf, h_kpts[k], s_kpts[k], overwrite, None,
symm_orb[k], irrep_id[k])
eig_kpts.append(e)
mo_coeff_kpts.append(c)
return eig_kpts, mo_coeff_kpts
[docs]
def eig_trs(kmf, h_kpts, s_kpts, overwrite=False, x=None):
''' Forcing real orbitals at time-reversal invariant momenta.
'''
cell = kmf.cell
kpts = kmf.kpts
trs_mask = is_trim(cell, kpts.kpts_ibz)
nao = h_kpts.shape[1]
nkpts = kpts.nkpts_ibz
eig_kpts = np.empty((nkpts, nao))
mo_coeff_kpts = np.empty((nkpts, nao, nao), dtype=h_kpts.dtype)
for k in range(nkpts):
h = h_kpts[k]
s = s_kpts[k]
if x is None:
if trs_mask[k]:
e, c = kmf._eigh(h.real, s.real)
else:
e, c = kmf._eigh(h, s)
eig_kpts[k] = e
mo_coeff_kpts[k] = c
else:
if trs_mask[k]:
e, c = kmf._eigh(h.real, s.real, x=x[k].real)
else:
e, c = kmf._eigh(h, s, x=x[k])
e1, _ = kmf._eigh(h, s)
nmo_k = c.shape[1]
eig_kpts[k,:nmo_k] = e
mo_coeff_kpts[k,:,:nmo_k] = c
if nmo_k < nao:
eig_kpts[k,nmo_k:] = pbchf.INVALID_ORBITAL_ENERGY
mo_coeff_kpts[k,:,nmo_k:] = 0
return eig_kpts, mo_coeff_kpts
[docs]
def ksymm_scf_common_init(kmf, cell, kpts, use_ao_symmetry=True):
kmf.use_ao_symmetry = (cell.dimension == 3 and
use_ao_symmetry and
not kpts.time_reversal and
kpts.symmorphic and
len(kpts.little_cogroup_ops) > 0)
if kmf.use_ao_symmetry and cell.symm_orb is None:
cell._build_symmetry(kpts)
return kmf
[docs]
class KsymAdaptedKSCF(khf.KSCF):
"""
KRHF with k-point symmetry
"""
_keys = {'use_ao_symmetry'}
get_occ = get_occ
get_rho = get_rho
energy_elec = energy_elec
def __init__(self, cell, kpts=libkpts.KPoints(),
exxdiv=getattr(__config__, 'pbc_scf_SCF_exxdiv', 'ewald'),
use_ao_symmetry=True):
ksymm_scf_common_init(self, cell, kpts, use_ao_symmetry)
khf.KSCF.__init__(self, cell, kpts=kpts, exxdiv=exxdiv)
@property
def kpts(self):
if 'kpts' in self.__dict__:
# To handle the attribute kpt loaded from chkfile
kpts_ibz = self.__dict__.pop('kpts')
if len(kpts_ibz) != self._kpts.nkpts_ibz:
raise RuntimeError("chkfile is not consistent with the current system.")
return self._kpts
@kpts.setter
def kpts(self, kpts):
if isinstance(kpts, np.ndarray):
logger.warn(self, "Input kpts is ndarray, building kpts object without symmetry.")
kpts = libkpts.make_kpts(self.cell, kpts=kpts)
elif not isinstance(kpts, libkpts.KPoints):
raise TypeError("Input kpts have wrong type: %s" % type(kpts))
self.with_df.kpts = kpts
self._kpts = kpts
@property
def kmesh(self):
from pyscf.pbc.tools.k2gamma import kpts_to_kmesh
kpts_bz = self._kpts.kpts
kmesh = kpts_to_kmesh(kpts_bz)
if len(kpts_bz) != np.prod(kmesh):
logger.WARN(self, 'K-points specified in %s are not Monkhorst-Pack %s grids',
self, kmesh)
return kmesh
@kmesh.setter
def kmesh(self, x):
self.kpts = self.cell.make_kpts(x, space_group_symmetry=True)
[docs]
def reset(self, cell=None):
self._kpts.reset(cell)
khf.KSCF.reset(self, cell)
return self
[docs]
def dump_flags(self, verbose=None):
mol_hf.SCF.dump_flags(self, verbose)
logger.info(self, '\n')
logger.info(self, '******** PBC SCF flags ********')
logger.info(self, 'N kpts (BZ) = %d', self.kpts.nkpts)
logger.debug(self, 'kpts (BZ) = %s', self.kpts.kpts)
logger.debug(self, 'kpts weights (BZ) = %s', self.kpts.weights)
logger.info(self, 'N kpts (IBZ) = %d', self.kpts.nkpts_ibz)
logger.debug(self, 'kpts (IBZ) = %s', self.kpts.kpts_ibz)
logger.debug(self, 'kpts weights (IBZ) = %s', self.kpts.weights_ibz)
logger.info(self, 'Exchange divergence treatment (exxdiv) = %s', self.exxdiv)
cell = self.cell
if ((cell.dimension >= 2 and cell.low_dim_ft_type != 'inf_vacuum') and
isinstance(self.exxdiv, str) and self.exxdiv.lower() == 'ewald'):
madelung = tools.pbc.madelung(cell, [self.kpts.kpts])
logger.info(self, ' madelung (= occupied orbital energy shift) = %s', madelung)
# FIXME: consider the fractional num_electron or not? This maybe
# relates to the charged system.
nelectron = float(self.cell.tot_electrons(self.kpts.nkpts)) / self.kpts.nkpts
logger.info(self, ' Total energy shift due to Ewald probe charge'
' = -1/2 * Nelec*madelung = %.12g',
madelung*nelectron * -.5)
logger.info(self, 'DF object = %s', self.with_df)
if not getattr(self.with_df, 'build', None):
self.with_df.dump_flags(verbose)
return self
[docs]
def check_linear_dependency(self, s, verbose=None):
cell = self.cell
symm_orb = cell.symm_orb
if symm_orb is None:
return khf.KSCF.check_linear_dependency(self, s, verbose)
assert len(symm_orb) == len(s), 'Number of k-points mismatch'
log = logger.new_logger(self, verbose)
irrep_id = cell.irrep_id
kpts = self.kpts
if isinstance(kpts, libkpts.KPoints):
kpts = kpts.kpts_ibz
trs_mask = is_trim(cell, kpts)
x_kpts = []
for k, s_k in enumerate(s):
nirrep = len(symm_orb[k])
if trs_mask[k]:
s_k = s_k.real
s_k = symmetrize_matrix(s_k, symm_orb[k])
orbsym = []
xs = []
for ir in range(nirrep):
e, v = scipy.linalg.eigh(s_k[ir])
if mol_hf.remove_overlap_zero_eigenvalue:
mask = e > mol_hf.overlap_zero_eigenvalue_threshold
x = v[:,mask] / np.sqrt(e[mask])
nao, nmo = x.shape
if nmo < nao:
log.info(f"kpt {k}: {nao-nmo} small eigenvectors of overlap matrix removed")
else:
x = v / np.sqrt(e)
xs.append(x)
orbsym.append(np.repeat(irrep_id[k][ir], x.shape[1]))
x_orth = so2ao_mo_coeff(symm_orb[k], xs)
x_kpts.append(lib.tag_array(x_orth, orbsym=np.hstack(orbsym)))
return x_kpts
[docs]
@lib.with_doc(khf.get_ovlp.__doc__)
def get_ovlp(self, cell=None, kpts=None):
if isinstance(kpts, np.ndarray):
return khf.KSCF.get_ovlp(self, cell, kpts)
if kpts is None: kpts = self.kpts
return khf.KSCF.get_ovlp(self, cell, kpts.kpts_ibz)
[docs]
@lib.with_doc(khf.get_hcore.__doc__)
def get_hcore(self, cell=None, kpts=None):
if isinstance(kpts, np.ndarray):
return khf.KSCF.get_hcore(self, cell, kpts)
if kpts is None: kpts = self.kpts
return khf.KSCF.get_hcore(self, cell, kpts.kpts_ibz)
[docs]
@lib.with_doc(khf.get_jk.__doc__)
def get_jk(self, cell=None, dm_kpts=None, hermi=1, kpts=None, kpts_band=None,
with_j=True, with_k=True, omega=None, **kwargs):
if isinstance(kpts, np.ndarray):
return super().get_jk(cell, dm_kpts, hermi, kpts, kpts_band,
with_j, with_k, omega, **kwargs)
if cell is None: cell = self.cell
if kpts is None: kpts = self.kpts
if dm_kpts is None: dm_kpts = self.make_rdm1()
#get dms for each kpt in BZ
if isinstance(dm_kpts[0], np.ndarray) and dm_kpts[0].ndim == 3:
ndm = len(dm_kpts[0])
else:
ndm = len(dm_kpts)
if ndm != kpts.nkpts_ibz:
raise RuntimeError("Number of input density matrices does not \
match the number of IBZ kpts: %d vs %d."
% (ndm, kpts.nkpts_ibz))
dm_kpts = kpts.transform_dm(dm_kpts)
if kpts_band is None: kpts_band = kpts.kpts_ibz
cpu0 = (logger.process_clock(), logger.perf_counter())
if self.rsjk:
raise NotImplementedError('rsjk with k-points symmetry')
else:
vj, vk = self.with_df.get_jk(dm_kpts, hermi, kpts.kpts, kpts_band,
with_j, with_k, omega, exxdiv=self.exxdiv)
logger.timer(self, 'vj and vk', *cpu0)
return vj, vk
[docs]
def init_guess_by_chkfile(self, chk=None, project=None, kpts=None):
if isinstance(kpts, np.ndarray):
return super().init_guess_by_chkfile(chk, project, kpts)
if kpts is None: kpts = self.kpts
return super().init_guess_by_chkfile(chk, project, kpts.kpts_ibz)
[docs]
def dump_chk(self, envs):
if self.chkfile:
mol_hf.SCF.dump_chk(self, envs)
with lib.H5FileWrap(self.chkfile, 'a') as fh5:
fh5['scf/kpts'] = self.kpts.kpts_ibz #FIXME Shall we rebuild kpts? If so, more info is needed.
return self
[docs]
def eig(self, h_kpts, s_kpts, overwrite=False, x=None):
if self.use_ao_symmetry:
return eig(self, h_kpts, s_kpts, overwrite, x)
elif self.kpts.time_reversal:
return eig_trs(self, h_kpts, s_kpts, overwrite, x)
else:
return khf.KSCF.eig(self, h_kpts, s_kpts, overwrite, x)
[docs]
def get_orbsym(self, mo_coeff=None, s=None):
if not self.use_ao_symmetry:
raise RuntimeError("AO symmetry not initiated")
from pyscf.scf.hf_symm import get_orbsym
if mo_coeff is None:
mo_coeff = self.mo_coeff
if s is None:
s = self.get_ovlp()
cell = self.cell
symm_orb = cell.symm_orb
irrep_id = cell.irrep_id
orbsym = []
for k in range(len(mo_coeff)):
orbsym_k = np.asarray(get_orbsym(cell, mo_coeff[k], s=s[k],
symm_orb=symm_orb[k], irrep_id=irrep_id[k]))
orbsym.append(orbsym_k)
return orbsym
orbsym = property(get_orbsym)
def _finalize(self):
khf.KSCF._finalize(self)
if not self.use_ao_symmetry:
return self
orbsym = self.get_orbsym()
for k, mo_e in enumerate(self.mo_energy):
idx = np.argsort(mo_e.round(9), kind='stable')
self.mo_energy[k] = self.mo_energy[k][idx]
self.mo_occ[k] = self.mo_occ[k][idx]
self.mo_coeff[k] = lib.tag_array(self.mo_coeff[k][:,idx], orbsym=orbsym[k][idx])
self.dump_chk({'e_tot': self.e_tot, 'mo_energy': self.mo_energy,
'mo_coeff': self.mo_coeff, 'mo_occ': self.mo_occ})
return self
[docs]
def to_khf(self):
'''transform to non-symmetry object
'''
from pyscf.pbc.scf import kuhf_ksymm, kghf_ksymm
from pyscf.pbc.scf import khf, kuhf, kghf
from pyscf.pbc.dft import krks, krks_ksymm, kuks, kuks_ksymm
from pyscf.scf import addons as mol_addons
def update_mo_(mf, mf1):
kpts = mf.kpts
if mf.mo_energy is not None:
mo_energy = kpts.transform_mo_energy(mf.mo_energy)
mo_occ = kpts.transform_mo_occ(mf.mo_occ)
if isinstance(mf, kghf_ksymm.KGHF):
mo_coeff = np.asarray(mf.mo_coeff)
nao = mo_coeff.shape[1] // 2
mo_coeff_alpha = kpts.transform_mo_coeff(mo_coeff[:,:nao])
mo_coeff_beta = kpts.transform_mo_coeff(mo_coeff[:,nao:])
mo_coeff = []
for k in range(len(mo_coeff_alpha)):
mo_coeff.append(np.vstack((mo_coeff_alpha[k], mo_coeff_beta[k])))
mo_coeff = np.asarray(mo_coeff)
else:
mo_coeff = kpts.transform_mo_coeff(mf.mo_coeff)
mf1.mo_coeff = mo_coeff
mf1.mo_occ = mo_occ
mf1.mo_energy = mo_energy
return mf1
known_cls = {KsymAdaptedKRHF : khf.KRHF,
kuhf_ksymm.KUHF : kuhf.KUHF,
kghf_ksymm.KGHF : kghf.KGHF,
krks_ksymm.KRKS : krks.KRKS,
kuks_ksymm.KUKS : kuks.KUKS}
out = mol_addons._object_without_soscf(self, known_cls, False)
out.__dict__.pop('kpts', None)
return update_mo_(self, out)
[docs]
def sfx2c1e(self):
raise NotImplementedError
x2c = x2c1e = sfx2c1e
[docs]
class KsymAdaptedKRHF(KsymAdaptedKSCF, khf.KRHF):
to_ks = khf.KRHF.to_ks
convert_from_ = khf.KRHF.convert_from_
[docs]
def get_init_guess(self, cell=None, key='minao', s1e=None):
if s1e is None:
s1e = self.get_ovlp(cell)
dm_kpts = mol_hf.SCF.get_init_guess(self, cell, key)
if dm_kpts.ndim == 2:
dm_kpts = np.asarray([dm_kpts]*self.kpts.nkpts_ibz)
elif len(dm_kpts) != self.kpts.nkpts_ibz:
dm_kpts = dm_kpts[self.kpts.ibz2bz]
ne = lib.einsum('k,kij,kji', self.kpts.weights_ibz, dm_kpts, s1e).real
nkpts = self.kpts.nkpts
ne *= nkpts
nelectron = float(self.cell.tot_electrons(nkpts))
if abs(ne - nelectron) > 0.01*nkpts:
logger.debug(self, 'Big error detected in the electron number '
'of initial guess density matrix (Ne/cell = %g)!\n'
' This can cause huge error in Fock matrix and '
'lead to instability in SCF for low-dimensional '
'systems.\n DM is normalized wrt the number '
'of electrons %s', ne/nkpts, nelectron/nkpts)
dm_kpts *= (nelectron / ne).reshape(-1,1,1)
return dm_kpts
KRHF = KsymAdaptedKRHF