def check_sanity(self):
lib.StreamObject.check_sanity(self)
cell = self.cell
if (cell.dimension < 2 or
(cell.dimension == 2 and cell.low_dim_ft_type == 'inf_vacuum')):
raise RuntimeError(
'FFTDF method does not support 0D/1D low-dimension '
'PBC system. DF, MDF or AFTDF methods should '
'be used.\nSee also examples/pbc/31-low_dimensional_pbc.py')
if not cell.has_ecp():
logger.warn(
self, 'FFTDF integrals are found in all-electron '
'calculation. It often causes huge error.\n'
'Recommended methods are DF or MDF. In SCF calculation, '
'they can be initialized as\n'
' mf = mf.density_fit()\nor\n'
' mf = mf.mix_density_fit()')
if cell.ke_cutoff is None:
ke_cutoff = tools.mesh_to_cutoff(cell.lattice_vectors(),
self.mesh).min()
else:
ke_cutoff = numpy.min(cell.ke_cutoff)
ke_guess = estimate_ke_cutoff(cell, cell.precision)
if ke_cutoff < ke_guess * KE_SCALING:
mesh_guess = tools.cutoff_to_mesh(cell.lattice_vectors(), ke_guess)
logger.warn(
self, 'ke_cutoff/mesh (%g / %s) is not enough for FFTDF '
'to get integral accuracy %g.\nCoulomb integral error '
'is ~ %.2g Eh.\nRecommended ke_cutoff/mesh are %g / %s.',
ke_cutoff, self.mesh, cell.precision,
error_for_ke_cutoff(cell, ke_cutoff), ke_guess, mesh_guess)
return self
def _mesh_for_valence(cell, valence_exp=VALENCE_EXP):
'''Energy cutoff estimation'''
b = cell.reciprocal_vectors()
if cell.dimension == 0:
w = 1
elif cell.dimension == 1:
w = numpy.linalg.norm(b[0]) / (2 * numpy.pi)
elif cell.dimension == 2:
w = numpy.linalg.norm(numpy.cross(b[0], b[1])) / (2 * numpy.pi)**2
else:
w = abs(numpy.linalg.det(b)) / (2 * numpy.pi)**3
precision = cell.precision * 10
Ecut_max = 0
for i in range(cell.nbas):
l = cell.bas_angular(i)
es = cell.bas_exp(i).copy()
es[es > valence_exp] = valence_exp
cs = abs(cell.bas_ctr_coeff(i)).max(axis=1)
ke_guess = gto.cell._estimate_ke_cutoff(es, l, cs, precision, w)
Ecut_max = max(Ecut_max, ke_guess.max())
mesh = tools.cutoff_to_mesh(cell.lattice_vectors(), Ecut_max)
mesh = numpy.min((mesh, cell.mesh), axis=0)
mesh[cell.dimension:] = cell.mesh[cell.dimension:]
return mesh
def __init__(self, cell, kpts=numpy.zeros((1, 3))):
self.cell = cell
self.stdout = cell.stdout
self.verbose = cell.verbose
self.max_memory = cell.max_memory
self.kpts = kpts # default is gamma point
self.kpts_band = None
self._auxbasis = None
# Search for optimized eta and mesh here.
if cell.dimension == 0:
self.eta = 0.2
self.mesh = cell.mesh
else:
ke_cutoff = tools.mesh_to_cutoff(cell.lattice_vectors(), cell.mesh)
ke_cutoff = ke_cutoff[:cell.dimension].min()
eta_cell = aft.estimate_eta_for_ke_cutoff(cell, ke_cutoff,
cell.precision)
eta_guess = estimate_eta(cell, cell.precision)
if eta_cell < eta_guess:
self.eta = eta_cell
# TODO? Round off mesh to the nearest odd numbers.
# Odd number of grids is preferred because even number of
# grids may break the conjugation symmetry between the
# k-points k and -k.
#?self.mesh = [(n//2)*2+1 for n in cell.mesh]
self.mesh = cell.mesh
else:
self.eta = eta_guess
ke_cutoff = aft.estimate_ke_cutoff_for_eta(
cell, self.eta, cell.precision)
self.mesh = tools.cutoff_to_mesh(cell.lattice_vectors(),
ke_cutoff)
if cell.dimension < 2 or cell.low_dim_ft_type == 'inf_vacuum':
self.mesh[cell.dimension:] = cell.mesh[cell.dimension:]
# exp_to_discard to remove diffused fitting functions. The diffused
# fitting functions may cause linear dependency in DF metric. Removing
# the fitting functions whose exponents are smaller than exp_to_discard
# can improve the linear dependency issue. However, this parameter
# affects the quality of the auxiliary basis. The default value of
# this parameter was set to 0.2 in v1.5.1 or older and was changed to
# 0 since v1.5.2.
self.exp_to_discard = cell.exp_to_discard
# The following attributes are not input options.
self.exxdiv = None # to mimic KRHF/KUHF object in function get_coulG
self.auxcell = None
self.blockdim = getattr(__config__, 'pbc_df_df_DF_blockdim', 240)
self.linear_dep_threshold = LINEAR_DEP_THR
self._j_only = False
# If _cderi_to_save is specified, the 3C-integral tensor will be saved in this file.
self._cderi_to_save = tempfile.NamedTemporaryFile(dir=lib.param.TMPDIR)
# If _cderi is specified, the 3C-integral tensor will be read from this file
self._cderi = None
self._keys = set(self.__dict__.keys())
def check_sanity(self):
lib.StreamObject.check_sanity(self)
cell = self.cell
if not cell.has_ecp():
logger.warn(
self, 'AFTDF integrals are found in all-electron '
'calculation. It often causes huge error.\n'
'Recommended methods are DF or MDF. In SCF calculation, '
'they can be initialized as\n'
' mf = mf.density_fit()\nor\n'
' mf = mf.mix_density_fit()')
if cell.dimension > 0:
if cell.ke_cutoff is None:
ke_cutoff = tools.mesh_to_cutoff(cell.lattice_vectors(),
self.mesh)
ke_cutoff = ke_cutoff[:cell.dimension].min()
else:
ke_cutoff = numpy.min(cell.ke_cutoff)
ke_guess = estimate_ke_cutoff(cell, cell.precision)
mesh_guess = tools.cutoff_to_mesh(cell.lattice_vectors(), ke_guess)
if ke_cutoff < ke_guess * KE_SCALING:
logger.warn(
self, 'ke_cutoff/mesh (%g / %s) is not enough for AFTDF '
'to get integral accuracy %g.\nCoulomb integral error '
'is ~ %.2g Eh.\nRecommended ke_cutoff/mesh are %g / %s.',
ke_cutoff, self.mesh, cell.precision,
error_for_ke_cutoff(cell, ke_cutoff), ke_guess, mesh_guess)
else:
mesh_guess = numpy.copy(self.mesh)
if cell.dimension < 3:
err = numpy.exp(-0.436392335 * min(self.mesh[cell.dimension:]) -
2.99944305)
err *= cell.nelectron
meshz = pbcgto.cell._mesh_inf_vaccum(cell)
mesh_guess[cell.dimension:] = int(meshz)
if err > cell.precision * 10:
logger.warn(
self, 'mesh %s of AFTDF may not be enough to get '
'integral accuracy %g for %dD PBC system.\n'
'Coulomb integral error is ~ %.2g Eh.\n'
'Recommended mesh is %s.', self.mesh, cell.precision,
cell.dimension, err, mesh_guess)
if (cell.mesh[cell.dimension:] / (1. * meshz) > 1.1).any():
meshz = pbcgto.cell._mesh_inf_vaccum(cell)
logger.warn(
self,
'setting mesh %s of AFTDF too high in non-periodic direction '
'(=%s) can result in an unnecessarily slow calculation.\n'
'For coulomb integral error of ~ %.2g Eh in %dD PBC, \n'
'a recommended mesh for non-periodic direction is %s.',
self.mesh, self.mesh[cell.dimension:], cell.precision,
cell.dimension, mesh_guess[cell.dimension:])
return self
def __init__(self, cell, kpts=numpy.zeros((1,3))):
self.cell = cell
self.stdout = cell.stdout
self.verbose = cell.verbose
self.max_memory = cell.max_memory
self.kpts = kpts # default is gamma point
self.kpts_band = None
self._auxbasis = None
# Search for optimized eta and mesh here.
if cell.dimension == 0:
self.eta = 0.2
self.mesh = cell.mesh
else:
ke_cutoff = tools.mesh_to_cutoff(cell.lattice_vectors(), cell.mesh)
ke_cutoff = ke_cutoff[:cell.dimension].min()
eta_cell = aft.estimate_eta_for_ke_cutoff(cell, ke_cutoff, cell.precision)
eta_guess = estimate_eta(cell, cell.precision)
if eta_cell < eta_guess:
self.eta = eta_cell
# TODO? Round off mesh to the nearest odd numbers.
# Odd number of grids is preferred because even number of
# grids may break the conjugation symmetry between the
# k-points k and -k.
#?self.mesh = [(n//2)*2+1 for n in cell.mesh]
self.mesh = cell.mesh
else:
self.eta = eta_guess
ke_cutoff = aft.estimate_ke_cutoff_for_eta(cell, self.eta, cell.precision)
self.mesh = tools.cutoff_to_mesh(cell.lattice_vectors(), ke_cutoff)
if cell.dimension < 2 or cell.low_dim_ft_type == 'inf_vacuum':
self.mesh[cell.dimension:] = cell.mesh[cell.dimension:]
# exp_to_discard to remove diffused fitting functions. The diffused
# fitting functions may cause linear dependency in DF metric. Removing
# the fitting functions whose exponents are smaller than exp_to_discard
# can improve the linear dependency issue. However, this parameter
# affects the quality of the auxiliary basis. The default value of
# this parameter was set to 0.2 in v1.5.1 or older and was changed to
# 0 since v1.5.2.
self.exp_to_discard = cell.exp_to_discard
# The following attributes are not input options.
self.exxdiv = None # to mimic KRHF/KUHF object in function get_coulG
self.auxcell = None
self.blockdim = getattr(__config__, 'pbc_df_df_DF_blockdim', 240)
self.linear_dep_threshold = LINEAR_DEP_THR
self._j_only = False
# If _cderi_to_save is specified, the 3C-integral tensor will be saved in this file.
self._cderi_to_save = tempfile.NamedTemporaryFile(dir=lib.param.TMPDIR)
# If _cderi is specified, the 3C-integral tensor will be read from this file
self._cderi = None
self._keys = set(self.__dict__.keys())
def get_jk(self,
dm,
hermi=1,
kpts=None,
kpts_band=None,
with_j=True,
with_k=True,
omega=None,
exxdiv=None):
if omega is not None: # J/K for RSH functionals
cell = self.cell
# * AFT is computationally more efficient than GDF if the Coulomb
# attenuation tends to the long-range role (i.e. small omega).
# * Note: changing to AFT integrator may cause small difference to
# the GDF integrator. If a very strict GDF result is desired,
# we can disable this trick by setting
# LONGRANGE_AFT_TURNOVER_THRESHOLD to 0.
# * The sparse mesh is not appropriate for low dimensional systems
# with infinity vacuum since the ERI may require large mesh to
# sample density in vacuum.
if (omega < LONGRANGE_AFT_TURNOVER_THRESHOLD
and cell.dimension >= 2
and cell.low_dim_ft_type != 'inf_vacuum'):
mydf = aft.AFTDF(cell, self.kpts)
mydf.ke_cutoff = aft.estimate_ke_cutoff_for_omega(cell, omega)
mydf.mesh = tools.cutoff_to_mesh(cell.lattice_vectors(),
mydf.ke_cutoff)
else:
mydf = self
return _sub_df_jk_(mydf, dm, hermi, kpts, kpts_band, with_j,
with_k, omega, exxdiv)
if kpts is None:
if numpy.all(self.kpts == 0):
# Gamma-point calculation by default
kpts = numpy.zeros(3)
else:
kpts = self.kpts
kpts = numpy.asarray(kpts)
if kpts.shape == (3, ):
return df_jk.get_jk(self, dm, hermi, kpts, kpts_band, with_j,
with_k, exxdiv)
vj = vk = None
if with_k:
vk = df_jk.get_k_kpts(self, dm, hermi, kpts, kpts_band, exxdiv)
if with_j:
vj = df_jk.get_j_kpts(self, dm, hermi, kpts, kpts_band)
return vj, vk
def check_sanity(self):
lib.StreamObject.check_sanity(self)
cell = self.cell
if not cell.has_ecp():
logger.warn(self, 'AFTDF integrals are found in all-electron '
'calculation. It often causes huge error.\n'
'Recommended methods are DF or MDF. In SCF calculation, '
'they can be initialized as\n'
' mf = mf.density_fit()\nor\n'
' mf = mf.mix_density_fit()')
if cell.dimension > 0:
if cell.ke_cutoff is None:
ke_cutoff = tools.mesh_to_cutoff(cell.lattice_vectors(), self.mesh)
ke_cutoff = ke_cutoff[:cell.dimension].min()
else:
ke_cutoff = numpy.min(cell.ke_cutoff)
ke_guess = estimate_ke_cutoff(cell, cell.precision)
mesh_guess = tools.cutoff_to_mesh(cell.lattice_vectors(), ke_guess)
if ke_cutoff < ke_guess * KE_SCALING:
logger.warn(self, 'ke_cutoff/mesh (%g / %s) is not enough for AFTDF '
'to get integral accuracy %g.\nCoulomb integral error '
'is ~ %.2g Eh.\nRecommended ke_cutoff/mesh are %g / %s.',
ke_cutoff, self.mesh, cell.precision,
error_for_ke_cutoff(cell, ke_cutoff), ke_guess, mesh_guess)
else:
mesh_guess = numpy.copy(self.mesh)
if cell.dimension < 3:
err = numpy.exp(-0.436392335*min(self.mesh[cell.dimension:]) - 2.99944305)
err *= cell.nelectron
meshz = pbcgto.cell._mesh_inf_vaccum(cell)
mesh_guess[cell.dimension:] = int(meshz)
if err > cell.precision*10:
logger.warn(self, 'mesh %s of AFTDF may not be enough to get '
'integral accuracy %g for %dD PBC system.\n'
'Coulomb integral error is ~ %.2g Eh.\n'
'Recommended mesh is %s.',
self.mesh, cell.precision, cell.dimension, err, mesh_guess)
if (cell.mesh[cell.dimension:]/(1.*meshz) > 1.1).any():
meshz = pbcgto.cell._mesh_inf_vaccum(cell)
logger.warn(self, 'setting mesh %s of AFTDF too high in non-periodic direction '
'(=%s) can result in an unnecessarily slow calculation.\n'
'For coulomb integral error of ~ %.2g Eh in %dD PBC, \n'
'a recommended mesh for non-periodic direction is %s.',
self.mesh, self.mesh[cell.dimension:], cell.precision,
cell.dimension, mesh_guess[cell.dimension:])
return self
def _mesh_for_valence(cell, valence_exp=VALENCE_EXP):
'''Energy cutoff estimation'''
precision = cell.precision * 10
Ecut_max = 0
for i in range(cell.nbas):
l = cell.bas_angular(i)
es = cell.bas_exp(i).copy()
es[es>valence_exp] = valence_exp
cs = abs(cell.bas_ctr_coeff(i)).max(axis=1)
ke_guess = gto.cell._estimate_ke_cutoff(es, l, cs, precision)
Ecut_max = max(Ecut_max, ke_guess.max())
mesh = tools.cutoff_to_mesh(cell.lattice_vectors(), Ecut_max)
mesh = numpy.min((mesh, cell.mesh), axis=0)
if cell.dimension < 2 or cell.low_dim_ft_type == 'inf_vacuum':
mesh[cell.dimension:] = cell.mesh[cell.dimension:]
return _round_off_to_odd_mesh(mesh)
def _mesh_for_valence(cell, valence_exp=VALENCE_EXP):
'''Energy cutoff estimation'''
precision = cell.precision * 10
Ecut_max = 0
for i in range(cell.nbas):
l = cell.bas_angular(i)
es = cell.bas_exp(i).copy()
es[es>valence_exp] = valence_exp
cs = abs(cell.bas_ctr_coeff(i)).max(axis=1)
ke_guess = gto.cell._estimate_ke_cutoff(es, l, cs, precision)
Ecut_max = max(Ecut_max, ke_guess.max())
mesh = tools.cutoff_to_mesh(cell.lattice_vectors(), Ecut_max)
mesh = numpy.min((mesh, cell.mesh), axis=0)
if cell.dimension < 2 or cell.low_dim_ft_type == 'inf_vacuum':
mesh[cell.dimension:] = cell.mesh[cell.dimension:]
return mesh
def __init__(self, cell, kpts=numpy.zeros((1, 3))):
self.cell = cell
self.stdout = cell.stdout
self.verbose = cell.verbose
self.max_memory = cell.max_memory
self.kpts = kpts # default is gamma point
self.kpts_band = None
self._auxbasis = None
if cell.dimension == 0:
self.eta = 0.2
self.mesh = cell.mesh
else:
ke_cutoff = tools.mesh_to_cutoff(cell.lattice_vectors(), cell.mesh)
ke_cutoff = ke_cutoff[:cell.dimension].min()
eta_cell = aft.estimate_eta_for_ke_cutoff(cell, ke_cutoff,
cell.precision)
eta_guess = estimate_eta(cell, cell.precision)
if eta_cell < eta_guess:
self.eta = eta_cell
self.mesh = cell.mesh
else:
self.eta = eta_guess
ke_cutoff = aft.estimate_ke_cutoff_for_eta(
cell, self.eta, cell.precision)
self.mesh = tools.cutoff_to_mesh(cell.lattice_vectors(),
ke_cutoff)
self.mesh[cell.dimension:] = cell.mesh[cell.dimension:]
# Not input options
self.exxdiv = None # to mimic KRHF/KUHF object in function get_coulG
self.auxcell = None
self.blockdim = getattr(__config__, 'pbc_df_df_DF_blockdim', 240)
self.linear_dep_threshold = LINEAR_DEP_THR
self._j_only = False
# If _cderi_to_save is specified, the 3C-integral tensor will be saved in this file.
self._cderi_to_save = tempfile.NamedTemporaryFile(dir=lib.param.TMPDIR)
# If _cderi is specified, the 3C-integral tensor will be read from this file
self._cderi = None
self._keys = set(self.__dict__.keys())
def get_pp_loc_part1(mydf, kpts=None):
cell = mydf.cell
if kpts is None:
kpts_lst = numpy.zeros((1, 3))
else:
kpts_lst = numpy.reshape(kpts, (-1, 3))
log = logger.Logger(mydf.stdout, mydf.verbose)
t0 = t1 = (time.clock(), time.time())
mesh = numpy.asarray(mydf.mesh)
nkpts = len(kpts_lst)
nao = cell.nao_nr()
nao_pair = nao * (nao + 1) // 2
charges = cell.atom_charges()
kpt_allow = numpy.zeros(3)
if mydf.eta == 0:
if cell.dimension > 0:
ke_guess = estimate_ke_cutoff(cell, cell.precision)
mesh_guess = tools.cutoff_to_mesh(cell.lattice_vectors(), ke_guess)
if numpy.any(
mesh[:cell.dimension] < mesh_guess[:cell.dimension] * .8):
logger.warn(
mydf, 'mesh %s is not enough for AFTDF.get_nuc function '
'to get integral accuracy %g.\nRecommended mesh is %s.',
mesh, cell.precision, mesh_guess)
Gv, Gvbase, kws = cell.get_Gv_weights(mesh)
vpplocG = pseudo.pp_int.get_gth_vlocG_part1(cell, Gv)
vpplocG = -numpy.einsum('ij,ij->j', cell.get_SI(Gv), vpplocG)
vpplocG *= kws
vG = vpplocG
vj = numpy.zeros((nkpts, nao_pair), dtype=numpy.complex128)
else:
if cell.dimension > 0:
ke_guess = estimate_ke_cutoff_for_eta(cell, mydf.eta,
cell.precision)
mesh_guess = tools.cutoff_to_mesh(cell.lattice_vectors(), ke_guess)
if numpy.any(mesh < mesh_guess * .8):
logger.warn(
mydf, 'mesh %s is not enough for AFTDF.get_nuc function '
'to get integral accuracy %g.\nRecommended mesh is %s.',
mesh, cell.precision, mesh_guess)
mesh_min = numpy.min((mesh_guess, mesh), axis=0)
if cell.dimension < 2 or cell.low_dim_ft_type == 'inf_vacuum':
mesh[:cell.dimension] = mesh_min[:cell.dimension]
else:
mesh = mesh_min
Gv, Gvbase, kws = cell.get_Gv_weights(mesh)
nuccell = _compensate_nuccell(mydf)
# PP-loc part1 is handled by fakenuc in _int_nuc_vloc
vj = lib.asarray(mydf._int_nuc_vloc(nuccell, kpts_lst))
t0 = t1 = log.timer_debug1('vnuc pass1: analytic int', *t0)
coulG = tools.get_coulG(cell, kpt_allow, mesh=mesh, Gv=Gv) * kws
aoaux = ft_ao.ft_ao(nuccell, Gv)
vG = numpy.einsum('i,xi->x', -charges, aoaux) * coulG
max_memory = max(2000, mydf.max_memory - lib.current_memory()[0])
for aoaoks, p0, p1 in mydf.ft_loop(mesh,
kpt_allow,
kpts_lst,
max_memory=max_memory,
aosym='s2'):
for k, aoao in enumerate(aoaoks):
# rho_ij(G) nuc(-G) / G^2
# = [Re(rho_ij(G)) + Im(rho_ij(G))*1j] [Re(nuc(G)) - Im(nuc(G))*1j] / G^2
if gamma_point(kpts_lst[k]):
vj[k] += numpy.einsum('k,kx->x', vG[p0:p1].real, aoao.real)
vj[k] += numpy.einsum('k,kx->x', vG[p0:p1].imag, aoao.imag)
else:
vj[k] += numpy.einsum('k,kx->x', vG[p0:p1].conj(), aoao)
t1 = log.timer_debug1('contracting Vnuc [%s:%s]' % (p0, p1), *t1)
log.timer_debug1('contracting Vnuc', *t0)
vj_kpts = []
for k, kpt in enumerate(kpts_lst):
if gamma_point(kpt):
vj_kpts.append(lib.unpack_tril(vj[k].real.copy()))
else:
vj_kpts.append(lib.unpack_tril(vj[k]))
if kpts is None or numpy.shape(kpts) == (3, ):
vj_kpts = vj_kpts[0]
return numpy.asarray(vj_kpts)
def get_pbc_pvxp(cell, kpts=None):
import numpy
import copy
import time
from pyscf import lib
from pyscf.lib import logger
from pyscf.pbc import tools
from pyscf.gto import mole
from pyscf.pbc.df import ft_ao
from pyscf.pbc.df import aft_jk
from pyscf.pbc.df import aft
if kpts is None:
kpts_lst = numpy.zeros((1,3))
else:
kpts_lst = numpy.reshape(kpts, (-1,3))
log = logger.Logger(cell.stdout, cell.verbose)
t1 = t0 = (time.clock(), time.time())
mydf = aft.AFTDF(cell, kpts)
mydf.eta = 0.2
ke_guess = aft.estimate_ke_cutoff_for_eta(cell, mydf.eta, cell.precision)
mydf.mesh = tools.cutoff_to_mesh(cell.lattice_vectors(), ke_guess)
log.debug('mydf.mesh %s', mydf.mesh)
nkpts = len(kpts_lst)
nao = cell.nao_nr()
nao_pair = nao * (nao+1) // 2
Gv, Gvbase, kws = cell.get_Gv_weights(mydf.mesh)
charge = -cell.atom_charges() # Apply Koseki effective charge?
kpt_allow = numpy.zeros(3)
coulG = tools.get_coulG(cell, kpt_allow, mesh=mydf.mesh, Gv=Gv)
coulG *= kws
if mydf.eta == 0:
soc_mat = numpy.zeros((nkpts,3,nao*nao), dtype=numpy.complex128)
SI = cell.get_SI(Gv)
vG = numpy.einsum('i,ix->x', charge, SI) * coulG
else:
nuccell = copy.copy(cell)
half_sph_norm = .5/numpy.sqrt(numpy.pi)
norm = half_sph_norm/mole.gaussian_int(2, mydf.eta)
chg_env = [mydf.eta, norm]
ptr_eta = cell._env.size
ptr_norm = ptr_eta + 1
chg_bas = [[ia, 0, 1, 1, 0, ptr_eta, ptr_norm, 0] for ia in range(cell.natm)]
nuccell._atm = cell._atm
nuccell._bas = numpy.asarray(chg_bas, dtype=numpy.int32)
nuccell._env = numpy.hstack((cell._env, chg_env))
soc_mat = mydf._int_nuc_vloc(nuccell, kpts_lst, 'int3c2e_pvxp1_sph',
aosym='s1', comp=3)
soc_mat = numpy.asarray(soc_mat).reshape(nkpts,3,nao**2)
t1 = log.timer_debug1('pnucp pass1: analytic int', *t1)
aoaux = ft_ao.ft_ao(nuccell, Gv)
vG = numpy.einsum('i,xi->x', charge, aoaux) * coulG
max_memory = max(2000, mydf.max_memory-lib.current_memory()[0])
for aoaoks, p0, p1 in mydf.ft_loop(mydf.mesh, kpt_allow, kpts_lst,
max_memory=max_memory, aosym='s1',
intor='GTO_ft_pxp_sph', comp=3):
for k, aoao in enumerate(aoaoks):
aoao = aoao.reshape(3,-1,nao**2)
if aft_jk.gamma_point(kpts_lst[k]):
soc_mat[k] += numpy.einsum('k,ckx->cx', vG[p0:p1].real, aoao.real)
soc_mat[k] += numpy.einsum('k,ckx->cx', vG[p0:p1].imag, aoao.imag)
else:
soc_mat[k] += numpy.einsum('k,ckx->cx', vG[p0:p1].conj(), aoao)
t1 = log.timer_debug1('contracting pnucp', *t1)
soc_mat_kpts = []
for k, kpt in enumerate(kpts_lst):
if aft_jk.gamma_point(kpt):
soc_mat_kpts.append(soc_mat[k].real.reshape(3,nao,nao))
else:
soc_mat_kpts.append(soc_mat[k].reshape(3,nao,nao))
if kpts is None or numpy.shape(kpts) == (3,):
soc_mat_kpts = soc_mat_kpts[0]
return numpy.asarray(soc_mat_kpts)
def build(self, omega=None, direct_scf_tol=None):
cpu0 = (time.clock(), time.time())
cell = self.cell
kpts = self.kpts
k_scaled = cell.get_scaled_kpts(kpts).sum(axis=0)
k_mod_to_half = k_scaled * 2 - (k_scaled * 2).round(0)
if abs(k_mod_to_half).sum() > 1e-5:
raise NotImplementedError('k-points must be symmetryic')
if omega is not None:
self.omega = omega
if self.omega is None:
# Search a proper range-separation parameter omega that can balance the
# computational cost between the real space integrals and moment space
# integrals
self.omega, self.mesh, self.ke_cutoff = _guess_omega(
cell, kpts, self.mesh)
else:
self.ke_cutoff = aft.estimate_ke_cutoff_for_omega(cell, self.omega)
self.mesh = pbctools.cutoff_to_mesh(cell.lattice_vectors(),
self.ke_cutoff)
logger.info(self, 'omega = %.15g ke_cutoff = %s mesh = %s',
self.omega, self.ke_cutoff, self.mesh)
if direct_scf_tol is None:
direct_scf_tol = cell.precision**1.5
logger.debug(self, 'Set direct_scf_tol %g', direct_scf_tol)
self.cell_rs = cell_rs = _re_contract_cell(cell, self.ke_cutoff)
self.bvk_kmesh = kmesh = k2gamma.kpts_to_kmesh(cell_rs, kpts)
bvkcell, phase = k2gamma.get_phase(cell_rs, kpts, kmesh)
self.bvkmesh_Ls = Ks = k2gamma.translation_vectors_for_kmesh(
cell_rs, kmesh)
self.bvkcell = bvkcell
self.phase = phase
# Given ke_cutoff, eta corresponds to the most steep Gaussian basis
# of which the Coulomb integrals can be accurately computed in moment
# space.
eta = aft.estimate_eta_for_ke_cutoff(cell,
self.ke_cutoff,
precision=cell.precision)
# * Assuming the most steep function in smooth basis has exponent eta,
# with attenuation parameter omega, rcut_sr is the distance of which
# the value of attenuated Coulomb integrals of four shells |eta> is
# smaller than the required precision.
# * The attenuated coulomb integrals between four s-type Gaussians
# (2*a/pi)^{3/4}exp(-a*r^2) is
# (erfc(omega*a^0.5/(omega^2+a)^0.5*R) - erfc(a^0.5*R)) / R
# if two Gaussians on one center and the other two on another center
# and the distance between the two centers are R.
# * The attenuated coulomb integrals between two spherical charge
# distributions is
# ~(pi/eta)^3/2 (erfc(tau*(eta/2)^0.5*R) - erfc((eta/2)^0.5*R)) / R
# tau = omega/sqrt(omega^2 + eta/2)
# if the spherical charge distribution is the product of above s-type
# Gaussian with exponent eta and a very smooth function.
# When R is large, the attenuated Coulomb integral is
# ~= (pi/eta)^3/2 erfc(tau*(eta/2)^0.5*R) / R
# ~= pi/(tau*eta^2*R^2) exp(-tau^2*eta*R^2/2)
tau = self.omega / (self.omega**2 + eta / 2)**.5
rcut_sr = 10 # initial guess
rcut_sr = (-np.log(direct_scf_tol * tau * (eta * rcut_sr)**2 / np.pi) /
(tau**2 * eta / 2))**.5
logger.debug(self, 'eta = %g rcut_sr = %g', eta, rcut_sr)
# Ls is the translation vectors to mimic periodicity of a cell
Ls = bvkcell.get_lattice_Ls(rcut=cell.rcut + rcut_sr)
self.supmol_Ls = Ls = Ls[np.linalg.norm(Ls, axis=1).argsort()]
supmol = _make_extended_mole(cell_rs, Ls, Ks, self.omega,
direct_scf_tol)
self.supmol = supmol
nkpts = len(self.bvkmesh_Ls)
nbas = cell_rs.nbas
n_steep, n_local, n_diffused = cell_rs._nbas_each_set
n_compact = n_steep + n_local
bas_mask = supmol._bas_mask
self.bvk_bas_mask = bvk_bas_mask = bas_mask.any(axis=2)
# Some basis in bvk-cell are not presented in the supmol. They can be
# skipped when computing SR integrals
self.bvkcell._bas = bvkcell._bas[bvk_bas_mask.ravel()]
# Record the mapping between the dense bvkcell basis and the
# original sparse bvkcell basis
bvk_cell_idx = np.repeat(np.arange(nkpts)[:, None], nbas, axis=1)
self.bvk_cell_id = bvk_cell_idx[bvk_bas_mask].astype(np.int32)
cell0_shl_idx = np.repeat(np.arange(nbas)[None, :], nkpts, axis=0)
self.cell0_shl_id = cell0_shl_idx[bvk_bas_mask].astype(np.int32)
logger.timer_debug1(self, 'initializing supmol', *cpu0)
logger.info(self, 'sup-mol nbas = %d cGTO = %d pGTO = %d', supmol.nbas,
supmol.nao, supmol.npgto_nr())
supmol.omega = -self.omega # Set short range coulomb
with supmol.with_integral_screen(direct_scf_tol**2):
vhfopt = _vhf.VHFOpt(supmol,
'int2e_sph',
qcondname=libpbc.PBCVHFsetnr_direct_scf)
vhfopt.direct_scf_tol = direct_scf_tol
self.vhfopt = vhfopt
logger.timer(self, 'initializing vhfopt', *cpu0)
q_cond = vhfopt.get_q_cond((supmol.nbas, supmol.nbas))
idx = supmol._images_loc
bvk_q_cond = lib.condense('NP_absmax', q_cond, idx, idx)
ovlp_mask = bvk_q_cond > direct_scf_tol
# Remove diffused-diffused block
if n_diffused > 0:
diffused_mask = np.zeros_like(bvk_bas_mask)
diffused_mask[:, n_compact:] = True
diffused_mask = diffused_mask[bvk_bas_mask]
ovlp_mask[diffused_mask[:, None] & diffused_mask] = False
self.ovlp_mask = ovlp_mask.astype(np.int8)
# mute rcut_threshold, divide basis into two sets only
cell_lr_aft = _re_contract_cell(cell, self.ke_cutoff, -1, verbose=0)
self.lr_aft = lr_aft = _LongRangeAFT(cell_lr_aft, kpts, self.omega,
self.bvk_kmesh)
lr_aft.ke_cutoff = self.ke_cutoff
lr_aft.mesh = self.mesh
lr_aft.eta = eta
return self
def get_nuc(mydf, kpts=None):
cell = mydf.cell
if kpts is None:
kpts_lst = numpy.zeros((1, 3))
else:
kpts_lst = numpy.reshape(kpts, (-1, 3))
log = logger.Logger(mydf.stdout, mydf.verbose)
t0 = t1 = (time.clock(), time.time())
mesh = numpy.asarray(mydf.mesh)
nkpts = len(kpts_lst)
nao = cell.nao_nr()
nao_pair = nao * (nao + 1) // 2
charges = cell.atom_charges()
kpt_allow = numpy.zeros(3)
if mydf.eta == 0:
if cell.dimension > 0:
ke_guess = estimate_ke_cutoff(cell, cell.precision)
mesh_guess = tools.cutoff_to_mesh(cell.lattice_vectors(), ke_guess)
if numpy.any(mesh < mesh_guess * .8):
logger.warn(
mydf, 'mesh %s is not enough for AFTDF.get_nuc function '
'to get integral accuracy %g.\nRecommended mesh is %s.',
mesh, cell.precision, mesh_guess)
Gv, Gvbase, kws = cell.get_Gv_weights(mesh)
vpplocG = pseudo.pp_int.get_gth_vlocG_part1(cell, Gv)
vpplocG = -numpy.einsum('ij,ij->j', cell.get_SI(Gv), vpplocG)
v1 = -vpplocG.copy()
if cell.dimension == 1 or cell.dimension == 2:
G0idx, SI_on_z = pbcgto.cell._SI_for_uniform_model_charge(cell, Gv)
coulG = 4 * numpy.pi / numpy.linalg.norm(Gv[G0idx], axis=1)**2
vpplocG[G0idx] += charges.sum() * SI_on_z * coulG
vpplocG *= kws
vG = vpplocG
vj = numpy.zeros((nkpts, nao_pair), dtype=numpy.complex128)
else:
if cell.dimension > 0:
ke_guess = estimate_ke_cutoff_for_eta(cell, mydf.eta,
cell.precision)
mesh_guess = tools.cutoff_to_mesh(cell.lattice_vectors(), ke_guess)
if numpy.any(mesh < mesh_guess * .8):
logger.warn(
mydf, 'mesh %s is not enough for AFTDF.get_nuc function '
'to get integral accuracy %g.\nRecommended mesh is %s.',
mesh, cell.precision, mesh_guess)
mesh_min = numpy.min(
(mesh_guess[:cell.dimension], mesh[:cell.dimension]), axis=0)
mesh[:cell.dimension] = mesh_min.astype(int)
Gv, Gvbase, kws = cell.get_Gv_weights(mesh)
nuccell = copy.copy(cell)
half_sph_norm = .5 / numpy.sqrt(numpy.pi)
norm = half_sph_norm / gto.gaussian_int(2, mydf.eta)
chg_env = [mydf.eta, norm]
ptr_eta = cell._env.size
ptr_norm = ptr_eta + 1
chg_bas = [[ia, 0, 1, 1, 0, ptr_eta, ptr_norm, 0]
for ia in range(cell.natm)]
nuccell._atm = cell._atm
nuccell._bas = numpy.asarray(chg_bas, dtype=numpy.int32)
nuccell._env = numpy.hstack((cell._env, chg_env))
# PP-loc part1 is handled by fakenuc in _int_nuc_vloc
vj = lib.asarray(mydf._int_nuc_vloc(nuccell, kpts_lst))
t0 = t1 = log.timer_debug1('vnuc pass1: analytic int', *t0)
coulG = tools.get_coulG(cell, kpt_allow, mesh=mesh, Gv=Gv) * kws
aoaux = ft_ao.ft_ao(nuccell, Gv)
vG = numpy.einsum('i,xi->x', -charges, aoaux) * coulG
if cell.dimension == 1 or cell.dimension == 2:
G0idx, SI_on_z = pbcgto.cell._SI_for_uniform_model_charge(cell, Gv)
vG[G0idx] += charges.sum() * SI_on_z * coulG[G0idx]
max_memory = max(2000, mydf.max_memory - lib.current_memory()[0])
for aoaoks, p0, p1 in mydf.ft_loop(mesh,
kpt_allow,
kpts_lst,
max_memory=max_memory,
aosym='s2'):
for k, aoao in enumerate(aoaoks):
# rho_ij(G) nuc(-G) / G^2
# = [Re(rho_ij(G)) + Im(rho_ij(G))*1j] [Re(nuc(G)) - Im(nuc(G))*1j] / G^2
if gamma_point(kpts_lst[k]):
vj[k] += numpy.einsum('k,kx->x', vG[p0:p1].real, aoao.real)
vj[k] += numpy.einsum('k,kx->x', vG[p0:p1].imag, aoao.imag)
else:
vj[k] += numpy.einsum('k,kx->x', vG[p0:p1].conj(), aoao)
t1 = log.timer_debug1('contracting Vnuc [%s:%s]' % (p0, p1), *t1)
log.timer_debug1('contracting Vnuc', *t0)
vj_kpts = []
for k, kpt in enumerate(kpts_lst):
if gamma_point(kpt):
vj_kpts.append(lib.unpack_tril(vj[k].real.copy()))
else:
vj_kpts.append(lib.unpack_tril(vj[k]))
if kpts is None or numpy.shape(kpts) == (3, ):
vj_kpts = vj_kpts[0]
return numpy.asarray(vj_kpts)
def get_pp_loc_part1(mydf, kpts=None):
cell = mydf.cell
if kpts is None:
kpts_lst = numpy.zeros((1,3))
else:
kpts_lst = numpy.reshape(kpts, (-1,3))
log = logger.Logger(mydf.stdout, mydf.verbose)
t0 = t1 = (time.clock(), time.time())
mesh = numpy.asarray(mydf.mesh)
nkpts = len(kpts_lst)
nao = cell.nao_nr()
nao_pair = nao * (nao+1) // 2
charges = cell.atom_charges()
kpt_allow = numpy.zeros(3)
if mydf.eta == 0:
if cell.dimension > 0:
ke_guess = estimate_ke_cutoff(cell, cell.precision)
mesh_guess = tools.cutoff_to_mesh(cell.lattice_vectors(), ke_guess)
if numpy.any(mesh[:cell.dimension] < mesh_guess[:cell.dimension]*.8):
logger.warn(mydf, 'mesh %s is not enough for AFTDF.get_nuc function '
'to get integral accuracy %g.\nRecommended mesh is %s.',
mesh, cell.precision, mesh_guess)
Gv, Gvbase, kws = cell.get_Gv_weights(mesh)
vpplocG = pseudo.pp_int.get_gth_vlocG_part1(cell, Gv)
vpplocG = -numpy.einsum('ij,ij->j', cell.get_SI(Gv), vpplocG)
vpplocG *= kws
vG = vpplocG
vj = numpy.zeros((nkpts,nao_pair), dtype=numpy.complex128)
else:
if cell.dimension > 0:
ke_guess = estimate_ke_cutoff_for_eta(cell, mydf.eta, cell.precision)
mesh_guess = tools.cutoff_to_mesh(cell.lattice_vectors(), ke_guess)
if numpy.any(mesh < mesh_guess*.8):
logger.warn(mydf, 'mesh %s is not enough for AFTDF.get_nuc function '
'to get integral accuracy %g.\nRecommended mesh is %s.',
mesh, cell.precision, mesh_guess)
mesh_min = numpy.min((mesh_guess, mesh), axis=0)
if cell.dimension < 2 or cell.low_dim_ft_type == 'inf_vacuum':
mesh[:cell.dimension] = mesh_min[:cell.dimension]
else:
mesh = mesh_min
Gv, Gvbase, kws = cell.get_Gv_weights(mesh)
nuccell = _compensate_nuccell(mydf)
# PP-loc part1 is handled by fakenuc in _int_nuc_vloc
vj = lib.asarray(mydf._int_nuc_vloc(nuccell, kpts_lst))
t0 = t1 = log.timer_debug1('vnuc pass1: analytic int', *t0)
coulG = tools.get_coulG(cell, kpt_allow, mesh=mesh, Gv=Gv) * kws
aoaux = ft_ao.ft_ao(nuccell, Gv)
vG = numpy.einsum('i,xi->x', -charges, aoaux) * coulG
max_memory = max(2000, mydf.max_memory-lib.current_memory()[0])
for aoaoks, p0, p1 in mydf.ft_loop(mesh, kpt_allow, kpts_lst,
max_memory=max_memory, aosym='s2'):
for k, aoao in enumerate(aoaoks):
# rho_ij(G) nuc(-G) / G^2
# = [Re(rho_ij(G)) + Im(rho_ij(G))*1j] [Re(nuc(G)) - Im(nuc(G))*1j] / G^2
if gamma_point(kpts_lst[k]):
vj[k] += numpy.einsum('k,kx->x', vG[p0:p1].real, aoao.real)
vj[k] += numpy.einsum('k,kx->x', vG[p0:p1].imag, aoao.imag)
else:
vj[k] += numpy.einsum('k,kx->x', vG[p0:p1].conj(), aoao)
t1 = log.timer_debug1('contracting Vnuc [%s:%s]'%(p0, p1), *t1)
log.timer_debug1('contracting Vnuc', *t0)
vj_kpts = []
for k, kpt in enumerate(kpts_lst):
if gamma_point(kpt):
vj_kpts.append(lib.unpack_tril(vj[k].real.copy()))
else:
vj_kpts.append(lib.unpack_tril(vj[k]))
if kpts is None or numpy.shape(kpts) == (3,):
vj_kpts = vj_kpts[0]
return numpy.asarray(vj_kpts)
def _dip_correction(mf):
'''Makov-Payne corrections for charged systems.'''
from pyscf.pbc import gto
from pyscf.pbc import tools
from pyscf.pbc.dft import gen_grid
log = logger.new_logger(mf)
cell = mf.cell
a = cell.lattice_vectors()
b = np.linalg.inv(a).T
grids = gen_grid.UniformGrids(cell)
ke_cutoff = gto.estimate_ke_cutoff(cell, 1e-5)
grids.mesh = tools.cutoff_to_mesh(a, ke_cutoff)
dm = mf.make_rdm1()
rho = mf.get_rho(dm, grids, mf.kpt)
origin = _search_dipole_gauge_origin(cell, grids, rho, log)
def shift_grids(r):
r_frac = lib.dot(r - origin, b.T)
# Grids on the boundary (r_frac == +/-0.5) of the new cell may lead to
# unbalanced contributions to the dipole moment. Exclude them from the
# dipole and quadrupole
r_frac[r_frac== 0.5] = 0
r_frac[r_frac==-0.5] = 0
r_frac[r_frac > 0.5] -= 1
r_frac[r_frac <-0.5] += 1
r = lib.dot(r_frac, a)
return r
# SC BCC FCC
madelung = (-2.83729747948, -3.63923344951, -4.58486207411)
vol = cell.vol
L = vol ** (1./3)
chg = cell.charge
# epsilon is the dielectric constant of the system. For systems
# surrounded by vacuum, epsilon = 1.
epsilon = 1
# Energy correction of point charges of a simple cubic lattice.
de_mono = - chg**2 * np.array(madelung) / (2 * L * epsilon)
# dipole energy correction
r_e = shift_grids(grids.coords)
r_nuc = shift_grids(cell.atom_coords())
charges = cell.atom_charges()
e_dip = np.einsum('g,g,gx->x', rho, grids.weights, r_e)
nuc_dip = np.einsum('g,gx->x', charges, r_nuc)
dip = nuc_dip - e_dip
de_dip = -2.*np.pi/(3*cell.vol) * np.linalg.norm(dip)**2
# quadrupole energy correction
if abs(a - np.eye(3)*L).max() > 1e-5:
logger.warn(mf, 'System is not cubic cell. Quadrupole energy '
'correction is inaccurate since it is developed based on '
'cubic cell.')
e_quad = np.einsum('g,g,gx,gx->', rho, grids.weights, r_e, r_e)
nuc_quad = np.einsum('g,gx,gx->', charges, r_nuc, r_nuc)
quad = nuc_quad - e_quad
de_quad = 2.*np.pi/(3*cell.vol) * quad
de = de_mono + de_dip + de_quad
return de_mono, de_dip, de_quad, de
def _dip_correction(mf):
'''Makov-Payne corrections for charged systems.'''
from pyscf.pbc import tools
from pyscf.pbc.dft import gen_grid
log = logger.new_logger(mf)
cell = mf.cell
a = cell.lattice_vectors()
b = np.linalg.inv(a).T
grids = gen_grid.UniformGrids(cell)
ke_cutoff = gto.estimate_ke_cutoff(cell, 1e-5)
grids.mesh = tools.cutoff_to_mesh(a, ke_cutoff)
dm = mf.make_rdm1()
rho = mf.get_rho(dm, grids, mf.kpt)
origin = _search_dipole_gauge_origin(cell, grids, rho, log)
def shift_grids(r):
r_frac = lib.dot(r - origin, b.T)
# Grids on the boundary (r_frac == +/-0.5) of the new cell may lead to
# unbalanced contributions to the dipole moment. Exclude them from the
# dipole and quadrupole
r_frac[r_frac == 0.5] = 0
r_frac[r_frac == -0.5] = 0
r_frac[r_frac > 0.5] -= 1
r_frac[r_frac < -0.5] += 1
r = lib.dot(r_frac, a)
return r
# SC BCC FCC
madelung = (-2.83729747948, -3.63923344951, -4.58486207411)
vol = cell.vol
L = vol**(1. / 3)
chg = cell.charge
# epsilon is the dielectric constant of the system. For systems
# surrounded by vacuum, epsilon = 1.
epsilon = 1
# Energy correction of point charges of a simple cubic lattice.
de_mono = -chg**2 * np.array(madelung) / (2 * L * epsilon)
# dipole energy correction
r_e = shift_grids(grids.coords)
r_nuc = shift_grids(cell.atom_coords())
charges = cell.atom_charges()
e_dip = np.einsum('g,g,gx->x', rho, grids.weights, r_e)
nuc_dip = np.einsum('g,gx->x', charges, r_nuc)
dip = nuc_dip - e_dip
de_dip = -2. * np.pi / (3 * cell.vol) * np.linalg.norm(dip)**2
# quadrupole energy correction
if abs(a - np.eye(3) * L).max() > 1e-5:
logger.warn(
mf, 'System is not cubic cell. Quadrupole energy '
'correction is inaccurate since it is developed based on '
'cubic cell.')
e_quad = np.einsum('g,g,gx,gx->', rho, grids.weights, r_e, r_e)
nuc_quad = np.einsum('g,gx,gx->', charges, r_nuc, r_nuc)
quad = nuc_quad - e_quad
de_quad = 2. * np.pi / (3 * cell.vol) * quad
de = de_mono + de_dip + de_quad
return de_mono, de_dip, de_quad, de
def multi_grids_tasks(cell, verbose=None):
log = lib.logger.new_logger(cell, verbose)
tasks = []
a = cell.lattice_vectors()
neighbour_images = lib.cartesian_prod(([0, -1, 1], [0, -1, 1], [0, -1, 1]))
# Remove the first one which is the unit cell itself
neighbour_images0 = neighbour_images
neighbour_images = neighbour_images[1:]
neighbour_images = neighbour_images.dot(a)
b = numpy.linalg.inv(a.T)
heights = 1. / numpy.linalg.norm(b, axis=1)
normal_vector = b * heights.reshape(-1, 1)
distance_to_edge = cell.atom_coords().dot(normal_vector.T)
#FIXME: if atoms out of unit cell
#assert(numpy.all(distance_to_edge >= 0))
distance_to_edge = numpy.hstack(
[distance_to_edge, heights - distance_to_edge])
min_distance_to_edge = distance_to_edge.min(axis=1)
# Split shells based on rcut
rcuts_pgto, kecuts_pgto = _primitive_gto_cutoff(cell)
ao_loc = cell.ao_loc_nr()
def make_cell_high_exp(shls_high, r0, r1):
cell_high = copy.copy(cell)
cell_high._bas = cell._bas.copy()
cell_high._env = cell._env.copy()
rcut_atom = [0] * cell.natm
ke_cutoff = 0
for ib in shls_high:
rc = rcuts_pgto[ib]
idx = numpy.where((r1 <= rc) & (rc < r0))[0]
np1 = len(idx)
cs = cell._libcint_ctr_coeff(ib)
np, nc = cs.shape
if np1 < np: # no pGTO splitting within the shell
pexp = cell._bas[ib, PTR_EXP]
pcoeff = cell._bas[ib, PTR_COEFF]
cs1 = cs[idx]
cell_high._env[pcoeff:pcoeff + cs1.size] = cs1.T.ravel()
cell_high._env[pexp:pexp + np1] = cell.bas_exp(ib)[idx]
cell_high._bas[ib, NPRIM_OF] = np1
ke_cutoff = max(ke_cutoff, kecuts_pgto[ib][idx].max())
ia = cell.bas_atom(ib)
rcut_atom[ia] = max(rcut_atom[ia], rc[idx].max())
cell_high._bas = cell_high._bas[shls_high]
ao_idx = numpy.hstack(
[numpy.arange(ao_loc[i], ao_loc[i + 1]) for i in shls_high])
return cell_high, ao_idx, ke_cutoff, rcut_atom
def make_cell_low_exp(shls_low, r0, r1):
cell_low = copy.copy(cell)
cell_low._bas = cell._bas.copy()
cell_low._env = cell._env.copy()
for ib in shls_low:
idx = numpy.where(r0 <= rcuts_pgto[ib])[0]
np1 = len(idx)
cs = cell._libcint_ctr_coeff(ib)
np, nc = cs.shape
if np1 < np: # no pGTO splitting within the shell
pexp = cell._bas[ib, PTR_EXP]
pcoeff = cell._bas[ib, PTR_COEFF]
cs1 = cs[idx]
cell_low._env[pcoeff:pcoeff + cs1.size] = cs1.T.ravel()
cell_low._env[pexp:pexp + np1] = cell.bas_exp(ib)[idx]
cell_low._bas[ib, NPRIM_OF] = np1
cell_low._bas = cell_low._bas[shls_low]
ao_idx = numpy.hstack(
[numpy.arange(ao_loc[i], ao_loc[i + 1]) for i in shls_low])
return cell_low, ao_idx
nao = cell.nao_nr()
rmax = a.max() * RMAX_FACTOR
n_delimeter = int(numpy.log(0.01 / rmax) / numpy.log(RMAX_RATIO))
rcut_delimeter = rmax * (RMAX_RATIO**numpy.arange(n_delimeter))
for r0, r1 in zip(numpy.append(1e9, rcut_delimeter),
numpy.append(rcut_delimeter, 0)):
# shells which have high exps (small rcut)
shls_high = [
ib for ib, rc in enumerate(rcuts_pgto)
if numpy.any((r1 <= rc) & (rc < r0))
]
if len(shls_high) == 0:
continue
cell_high, ao_idx_high, ke_cutoff, rcut_atom = \
make_cell_high_exp(shls_high, r0, r1)
# shells which have low exps (big rcut)
shls_low = [
ib for ib, rc in enumerate(rcuts_pgto) if numpy.any(r0 <= rc)
]
if len(shls_low) == 0:
cell_low = None
ao_idx_low = []
else:
cell_low, ao_idx_low = make_cell_low_exp(shls_low, r0, r1)
mesh = tools.cutoff_to_mesh(a, ke_cutoff)
if TO_EVEN_GRIDS:
mesh = (mesh + 1) // 2 * 2 # to the nearest even number
if numpy.all(mesh >= cell.mesh):
# Including all rest shells
shls_high = [
ib for ib, rc in enumerate(rcuts_pgto) if numpy.any(rc < r0)
]
cell_high, ao_idx_high = make_cell_high_exp(shls_high, r0, 0)[:2]
cell_high.mesh = mesh = numpy.min([mesh, cell.mesh], axis=0)
coords = cell.gen_uniform_grids(mesh)
coords_f4 = coords.astype(numpy.float32)
ngrids = coords_f4.shape[0]
coords_idx = numpy.zeros(ngrids, dtype=bool)
gg = numpy.einsum('px,px->p', coords_f4, coords_f4)
Lg = (2 * neighbour_images.astype(numpy.float32)).dot(coords_f4.T)
for ia in set(cell_high._bas[:, ATOM_OF]):
rcut = rcut_atom[ia]
log.debug1(' atom %d rcut %g', mesh, ia, rcut)
atom_coord = cell.atom_coord(ia)
#dr = coords_f4 - atom_coord.astype(numpy.float32)
#coords_idx |= numpy.einsum('px,px->p', dr, dr) <= rcut**2
#optimized to
gg_ag = gg - coords_f4.dot(2 * atom_coord.astype(numpy.float32))
coords_idx |= gg_ag <= rcut**2 - atom_coord.dot(atom_coord)
if min_distance_to_edge[ia] > rcut:
# atom + rcut is completely inside the unit cell
continue
atoms_in_neighbour = neighbour_images + atom_coord
distance_to_unit_cell = atoms_in_neighbour.dot(normal_vector.T)
distance_to_unit_cell = numpy.hstack([
abs(distance_to_unit_cell),
abs(heights - distance_to_unit_cell)
])
idx = distance_to_unit_cell.min(axis=1) <= rcut
#for r_atom in atoms_in_neighbour[idx]:
# dr = coords_f4 - r_atom.astype(numpy.float32)
# coords_idx |= numpy.einsum('px,px->p', dr, dr) <= rcut**2
#optimized to
for i in numpy.where(idx)[0]:
L_a = atoms_in_neighbour[i]
coords_idx |= gg_ag - Lg[i] <= rcut**2 - L_a.dot(L_a)
coords = coords[coords_idx]
grids_high = gen_grid.UniformGrids(cell_high)
grids_high.coords = coords
grids_high.non0tab = grids_high.make_mask(cell_high, coords)
grids_high.coords_idx = coords_idx
grids_high.ao_idx = ao_idx_high
if cell_low is None:
grids_low = None
else:
grids_low = gen_grid.UniformGrids(cell_low)
grids_low.coords = coords
grids_low.non0tab = grids_low.make_mask(cell_low, coords)
grids_low.coords_idx = coords_idx
grids_low.ao_idx = ao_idx_low
log.debug('mesh %s nao all/high/low %d %d %d, ngrids %d', mesh, nao,
len(ao_idx_high), len(ao_idx_low), coords.shape[0])
tasks.append([grids_high, grids_low])
if numpy.all(mesh >= cell.mesh):
break
return tasks