def jk_part(mol, grid_coords, dms, fg):
# transfer bvv to SGXsetnr_direct_scf_blk. from _vhf.VHFOpt
# need add mol._bvv in scf.mole.py
c_bvv = numpy.asarray(mol._bvv, dtype=numpy.int32, order='C')
nbvv = ctypes.c_int(c_bvv.shape[0])
ao_loc = make_loc(c_bas, intor)
fsetqcond = getattr(libcvhf, 'SGXsetnr_direct_scf_blk')
fsetqcond(vhfopt._this, getattr(libcvhf, intor), lib.c_null_ptr(),
ao_loc.ctypes.data_as(ctypes.c_void_p),
c_atm.ctypes.data_as(ctypes.c_void_p), natm,
c_bas.ctypes.data_as(ctypes.c_void_p), nbas,
c_env.ctypes.data_as(ctypes.c_void_p),
c_bvv.ctypes.data_as(ctypes.c_void_p), nbvv)
fakemol = gto.fakemol_for_charges(grid_coords)
atm, bas, env = gto.mole.conc_env(mol._atm, mol._bas, mol._env,
fakemol._atm, fakemol._bas,
fakemol._env)
ao_loc = moleintor.make_loc(bas, intor)
shls_slice = (0, mol.nbas, 0, mol.nbas, mol.nbas, len(bas))
ngrids = grid_coords.shape[0]
vj = vk = None
fjk = []
dmsptr = []
vjkptr = []
if with_j:
if dms[0].ndim == 1: # the value of density at each grid
vj = numpy.zeros((len(dms), ncomp, nao, nao))[:, 0]
for i, dm in enumerate(dms):
dmsptr.append(dm.ctypes.data_as(ctypes.c_void_p))
vjkptr.append(vj[i].ctypes.data_as(ctypes.c_void_p))
fjk.append(_vhf._fpointer('SGXnr' + aosym + '_ijg_g_ij'))
else:
vj = numpy.zeros((len(dms), ncomp, ngrids))[:, 0]
for i, dm in enumerate(dms):
dmsptr.append(dm.ctypes.data_as(ctypes.c_void_p))
vjkptr.append(vj[i].ctypes.data_as(ctypes.c_void_p))
fjk.append(_vhf._fpointer('SGXnr' + aosym + '_ijg_ji_g'))
if with_k:
vk = numpy.zeros((len(fg), ncomp, ngrids, nao))[:, 0]
for i, dm in enumerate(fg):
dmsptr.append(dm.ctypes.data_as(ctypes.c_void_p))
vjkptr.append(vk[i].ctypes.data_as(ctypes.c_void_p))
fjk.append(_vhf._fpointer('SGXnr' + aosym + '_ijg_gj_gi'))
n_dm = len(fjk)
fjk = (ctypes.c_void_p * (n_dm))(*fjk)
dmsptr = (ctypes.c_void_p * (n_dm))(*dmsptr)
vjkptr = (ctypes.c_void_p * (n_dm))(*vjkptr)
drv(cintor, fdot, fjk, dmsptr, vjkptr, n_dm, ncomp,
(ctypes.c_int * 6)(*shls_slice),
ao_loc.ctypes.data_as(ctypes.c_void_p), cintopt, vhfopt._this,
atm.ctypes.data_as(ctypes.c_void_p), ctypes.c_int(mol.natm),
bas.ctypes.data_as(ctypes.c_void_p), ctypes.c_int(mol.nbas),
env.ctypes.data_as(ctypes.c_void_p))
return vj, vk
def __init__(self, mol, intor, prescreen='CVHFnoscreen', qcondname=None):
intor = ascint3(intor)
self._this = ctypes.POINTER(_vhf._CVHFOpt)()
#print self._this.contents, expect ValueError: NULL pointer access
self._intor = intor
c_atm = numpy.asarray(mol._atm, dtype=numpy.int32, order='C')
c_bas = numpy.asarray(mol._bas, dtype=numpy.int32, order='C')
c_env = numpy.asarray(mol._env, dtype=numpy.double, order='C')
natm = ctypes.c_int(c_atm.shape[0])
nbas = ctypes.c_int(c_bas.shape[0])
self._cintopt = make_cintopt(c_atm, c_bas, c_env, intor)
libao2mo.CVHFinit_optimizer(ctypes.byref(self._this),
c_atm.ctypes.data_as(ctypes.c_void_p),
natm,
c_bas.ctypes.data_as(ctypes.c_void_p),
nbas,
c_env.ctypes.data_as(ctypes.c_void_p))
self._this.contents.fprescreen = _fpointer(prescreen)
if prescreen != 'CVHFnoscreen' and intor in ('int2e_sph',
'int2e_cart'):
# for int2e_sph, qcondname is 'CVHFsetnr_direct_scf'
ao_loc = make_loc(c_bas, self._intor)
fsetqcond = getattr(libao2mo, qcondname)
fsetqcond(self._this, getattr(libao2mo, intor), self._cintopt,
ao_loc.ctypes.data_as(ctypes.c_void_p),
c_atm.ctypes.data_as(ctypes.c_void_p), natm,
c_bas.ctypes.data_as(ctypes.c_void_p), nbas,
c_env.ctypes.data_as(ctypes.c_void_p))
def __init__(self, mol, intor, prescreen='CVHFnoscreen', qcondname=None):
intor = ascint3(intor)
self._this = ctypes.POINTER(_vhf._CVHFOpt)()
#print self._this.contents, expect ValueError: NULL pointer access
self._intor = intor
c_atm = numpy.asarray(mol._atm, dtype=numpy.int32, order='C')
c_bas = numpy.asarray(mol._bas, dtype=numpy.int32, order='C')
c_env = numpy.asarray(mol._env, dtype=numpy.double, order='C')
natm = ctypes.c_int(c_atm.shape[0])
nbas = ctypes.c_int(c_bas.shape[0])
self._cintopt = make_cintopt(c_atm, c_bas, c_env, intor)
libao2mo.CVHFinit_optimizer(ctypes.byref(self._this),
c_atm.ctypes.data_as(ctypes.c_void_p), natm,
c_bas.ctypes.data_as(ctypes.c_void_p), nbas,
c_env.ctypes.data_as(ctypes.c_void_p))
self._this.contents.fprescreen = _fpointer(prescreen)
if prescreen != 'CVHFnoscreen' and intor in ('int2e_sph', 'int2e_cart'):
# for int2e_sph, qcondname is 'CVHFsetnr_direct_scf'
ao_loc = make_loc(c_bas, intor)
fsetqcond = getattr(libao2mo, qcondname)
fsetqcond(self._this,
getattr(libao2mo, intor), self._cintopt,
ao_loc.ctypes.data_as(ctypes.c_void_p),
c_atm.ctypes.data_as(ctypes.c_void_p), natm,
c_bas.ctypes.data_as(ctypes.c_void_p), nbas,
c_env.ctypes.data_as(ctypes.c_void_p))
def set_dm(self, dm, atm, bas, env):
if self._dmcondname is not None:
c_atm = numpy.asarray(atm, dtype=numpy.int32, order='C')
c_bas = numpy.asarray(bas, dtype=numpy.int32, order='C')
c_env = numpy.asarray(env, dtype=numpy.double, order='C')
natm = ctypes.c_int(c_atm.shape[0])
nbas = ctypes.c_int(c_bas.shape[0])
if isinstance(dm, numpy.ndarray) and dm.ndim == 2:
n_dm = 1
ngrids = dm.shape[0]
else:
n_dm = len(dm)
ngrids = dm.shape[1]
dm = numpy.asarray(dm, order='C')
ao_loc = make_loc(c_bas, self._intor)
if isinstance(self._dmcondname, ctypes._CFuncPtr):
fsetdm = self._dmcondname
else:
fsetdm = getattr(libcvhf, self._dmcondname)
if self._dmcondname == 'SGXsetnr_direct_scf_dm':
fsetdm(self._this, dm.ctypes.data_as(ctypes.c_void_p),
ctypes.c_int(n_dm),
ao_loc.ctypes.data_as(ctypes.c_void_p),
c_atm.ctypes.data_as(ctypes.c_void_p), natm,
c_bas.ctypes.data_as(ctypes.c_void_p), nbas,
c_env.ctypes.data_as(ctypes.c_void_p),
ctypes.c_int(ngrids))
self.ngrids = ngrids
else:
raise ValueError('Can only use SGX dm screening for SGXOpt')
def r_e1(intor, mo_coeff, orbs_slice, sh_range, atm, bas, env,
tao, aosym='s1', comp=1, ao2mopt=None, out=None):
assert(aosym in ('s4', 's2ij', 's2kl', 's1', 'a2ij', 'a2kl', 'a4ij',
'a4kl', 'a4'))
intor = ascint3(intor)
mo_coeff = numpy.asfortranarray(mo_coeff)
i0, i1, j0, j1 = orbs_slice
icount = i1 - i0
jcount = j1 - j0
ij_count = icount * jcount
c_atm = numpy.asarray(atm, dtype=numpy.int32)
c_bas = numpy.asarray(bas, dtype=numpy.int32)
c_env = numpy.asarray(env)
natm = ctypes.c_int(c_atm.shape[0])
nbas = ctypes.c_int(c_bas.shape[0])
klsh0, klsh1, nkl = sh_range
if icount <= jcount:
fmmm = _fpointer('AO2MOmmm_r_iltj')
else:
fmmm = _fpointer('AO2MOmmm_r_igtj')
if out is None:
out = numpy.empty((comp,nkl,ij_count), dtype=numpy.complex)
else:
out = numpy.ndarray((comp,nkl,nao_pair), dtype=numpy.complex,
buffer=out)
if out.size == 0:
return out
if ao2mopt is not None:
cao2mopt = ao2mopt._this
cintopt = ao2mopt._cintopt
cintor = ao2mopt._intor
else:
cao2mopt = lib.c_null_ptr()
cintor = _fpointer(intor)
cintopt = make_cintopt(c_atm, c_bas, c_env, intor)
tao = numpy.asarray(tao, dtype=numpy.int32)
ao_loc = make_loc(bas, 'spinor')
fdrv = getattr(libao2mo, 'AO2MOr_e1_drv')
fill = _fpointer('AO2MOfill_r_' + aosym)
ftrans = _fpointer('AO2MOtranse1_r_' + aosym)
fdrv(cintor, fill, ftrans, fmmm,
out.ctypes.data_as(ctypes.c_void_p),
mo_coeff.ctypes.data_as(ctypes.c_void_p),
ctypes.c_int(klsh0), ctypes.c_int(klsh1-klsh0),
ctypes.c_int(nkl), ctypes.c_int(comp),
(ctypes.c_int*4)(*orbs_slice), tao.ctypes.data_as(ctypes.c_void_p),
ao_loc.ctypes.data_as(ctypes.c_void_p), cintopt, cao2mopt,
c_atm.ctypes.data_as(ctypes.c_void_p), natm,
c_bas.ctypes.data_as(ctypes.c_void_p), nbas,
c_env.ctypes.data_as(ctypes.c_void_p))
return out
def r_e1(intor, mo_coeff, orbs_slice, sh_range, atm, bas, env,
tao, aosym='s1', comp=1, ao2mopt=None, out=None):
assert(aosym in ('s4', 's2ij', 's2kl', 's1', 'a2ij', 'a2kl', 'a4ij',
'a4kl', 'a4'))
mo_coeff = numpy.asfortranarray(mo_coeff)
i0, i1, j0, j1 = orbs_slice
icount = i1 - i0
jcount = j1 - j0
ij_count = icount * jcount
c_atm = numpy.asarray(atm, dtype=numpy.int32)
c_bas = numpy.asarray(bas, dtype=numpy.int32)
c_env = numpy.asarray(env)
natm = ctypes.c_int(c_atm.shape[0])
nbas = ctypes.c_int(c_bas.shape[0])
klsh0, klsh1, nkl = sh_range
if icount <= jcount:
fmmm = _fpointer('AO2MOmmm_r_iltj')
else:
fmmm = _fpointer('AO2MOmmm_r_igtj')
if out is None:
out = numpy.empty((comp,nkl,ij_count), dtype=numpy.complex)
else:
out = numpy.ndarray((comp,nkl,nao_pair), dtype=numpy.complex,
buffer=out)
if out.size == 0:
return out
if ao2mopt is not None:
cao2mopt = ao2mopt._this
cintopt = ao2mopt._cintopt
cintor = ao2mopt._intor
else:
cao2mopt = lib.c_null_ptr()
cintor = _fpointer(intor)
cintopt = _vhf.make_cintopt(c_atm, c_bas, c_env, intor)
tao = numpy.asarray(tao, dtype=numpy.int32)
ao_loc = moleintor.make_loc(bas, 'spinor')
fdrv = getattr(libao2mo, 'AO2MOr_e1_drv')
fill = _fpointer('AO2MOfill_r_' + aosym)
ftrans = _fpointer('AO2MOtranse1_r_' + aosym)
fdrv(cintor, fill, ftrans, fmmm,
out.ctypes.data_as(ctypes.c_void_p),
mo_coeff.ctypes.data_as(ctypes.c_void_p),
ctypes.c_int(klsh0), ctypes.c_int(klsh1-klsh0),
ctypes.c_int(nkl), ctypes.c_int(comp),
(ctypes.c_int*4)(*orbs_slice), tao.ctypes.data_as(ctypes.c_void_p),
ao_loc.ctypes.data_as(ctypes.c_void_p), cintopt, cao2mopt,
c_atm.ctypes.data_as(ctypes.c_void_p), natm,
c_bas.ctypes.data_as(ctypes.c_void_p), nbas,
c_env.ctypes.data_as(ctypes.c_void_p))
return out
def jk_part(mol, grid_coords, dms, fg, weights):
atm, bas, env = mol._atm, mol._bas, mol._env
ngrids = grid_coords.shape[0]
env = numpy.append(env, grid_coords.ravel())
env[gto.NGRIDS] = ngrids
env[gto.PTR_GRIDS] = mol._env.size
if pjs:
sgxopt.set_dm(fg / numpy.sqrt(numpy.abs(weights[None,:])),
mol._atm, mol._bas, env)
ao_loc = moleintor.make_loc(bas, sgxopt._intor)
shls_slice = (0, mol.nbas, 0, mol.nbas)
fg = numpy.ascontiguousarray(fg.transpose(0,2,1))
vj = vk = None
fjk = []
dmsptr = []
vjkptr = []
if with_j:
if dms[0].ndim == 1: # the value of density at each grid
vj = numpy.zeros((len(dms),ncomp,nao,nao))[:,0]
for i, dm in enumerate(dms):
dmsptr.append(dm.ctypes.data_as(ctypes.c_void_p))
vjkptr.append(vj[i].ctypes.data_as(ctypes.c_void_p))
fjk.append(_vhf._fpointer('SGXnr'+aosym+'_ijg_g_ij'))
else:
vj = numpy.zeros((len(dms),ncomp,ngrids))[:,0]
for i, dm in enumerate(dms):
dmsptr.append(dm.ctypes.data_as(ctypes.c_void_p))
vjkptr.append(vj[i].ctypes.data_as(ctypes.c_void_p))
fjk.append(_vhf._fpointer('SGXnr'+aosym+'_ijg_ji_g'))
if with_k:
vk = numpy.zeros((len(fg),ncomp,nao,ngrids))[:,0]
for i, dm in enumerate(fg):
dmsptr.append(dm.ctypes.data_as(ctypes.c_void_p))
vjkptr.append(vk[i].ctypes.data_as(ctypes.c_void_p))
fjk.append(_vhf._fpointer('SGXnr'+aosym+'_ijg_gj_gi'))
n_dm = len(fjk)
fjk = (ctypes.c_void_p*(n_dm))(*fjk)
dmsptr = (ctypes.c_void_p*(n_dm))(*dmsptr)
vjkptr = (ctypes.c_void_p*(n_dm))(*vjkptr)
drv(cintor, fdot, fjk, dmsptr, vjkptr, n_dm, ncomp,
(ctypes.c_int*4)(*shls_slice),
ao_loc.ctypes.data_as(ctypes.c_void_p),
sgxopt._cintopt, sgxopt._this,
atm.ctypes.data_as(ctypes.c_void_p), ctypes.c_int(mol.natm),
bas.ctypes.data_as(ctypes.c_void_p), ctypes.c_int(mol.nbas),
env.ctypes.data_as(ctypes.c_void_p),
ctypes.c_int(env.shape[0]),
ctypes.c_int(2 if aosym == 's2' else 1))
if vk is not None:
vk = vk.transpose(0,2,1)
vk = numpy.ascontiguousarray(vk)
return vj, vk
def nr_e1fill(intor,
sh_range,
atm,
bas,
env,
aosym='s1',
comp=1,
ao2mopt=None,
out=None):
assert (aosym in ('s4', 's2ij', 's2kl', 's1'))
c_atm = numpy.asarray(atm, dtype=numpy.int32, order='C')
c_bas = numpy.asarray(bas, dtype=numpy.int32, order='C')
c_env = numpy.asarray(env, order='C')
natm = ctypes.c_int(c_atm.shape[0])
nbas = ctypes.c_int(c_bas.shape[0])
ao_loc = moleintor.make_loc(bas, intor)
nao = ao_loc[-1]
klsh0, klsh1, nkl = sh_range
if aosym in ('s4', 's2ij'):
nao_pair = nao * (nao + 1) // 2
else:
nao_pair = nao * nao
if out is None:
out = numpy.empty((comp, nkl, nao_pair))
else:
out = numpy.ndarray((comp, nkl, nao_pair), buffer=out)
if out.size == 0:
return out
if ao2mopt is not None:
cao2mopt = ao2mopt._this
cintopt = ao2mopt._cintopt
cintor = ao2mopt._intor
else:
cao2mopt = lib.c_null_ptr()
cintor = _fpointer(intor)
cintopt = _vhf.make_cintopt(c_atm, c_bas, c_env, intor)
fdrv = getattr(libao2mo, 'AO2MOnr_e1fill_drv')
fill = _fpointer('AO2MOfill_nr_' + aosym)
fdrv(cintor, fill,
out.ctypes.data_as(ctypes.c_void_p), ctypes.c_int(klsh0),
ctypes.c_int(klsh1 - klsh0), ctypes.c_int(nkl), ctypes.c_int(comp),
ao_loc.ctypes.data_as(ctypes.c_void_p), cintopt, cao2mopt,
c_atm.ctypes.data_as(ctypes.c_void_p), natm,
c_bas.ctypes.data_as(ctypes.c_void_p), nbas,
c_env.ctypes.data_as(ctypes.c_void_p))
return out
def init_cvhf_direct(self, mol, intor, qcondname):
intor = mol._add_suffix(intor)
if intor == self._intor:
cintopt = self._cintopt
else:
cintopt = lib.c_null_ptr()
ao_loc = make_loc(mol._bas, intor)
fsetqcond = getattr(libcvhf, qcondname)
natm = ctypes.c_int(mol.natm)
nbas = ctypes.c_int(mol.nbas)
fsetqcond(self._this, getattr(libcvhf, intor), cintopt,
ao_loc.ctypes.data_as(ctypes.c_void_p),
mol._atm.ctypes.data_as(ctypes.c_void_p), natm,
mol._bas.ctypes.data_as(ctypes.c_void_p), nbas,
mol._env.ctypes.data_as(ctypes.c_void_p))
def jk_part(mol, grid_coords, dms, fg):
fakemol = gto.fakemol_for_charges(grid_coords)
atm, bas, env = gto.mole.conc_env(mol._atm, mol._bas, mol._env,
fakemol._atm, fakemol._bas, fakemol._env)
ao_loc = moleintor.make_loc(bas, sgxopt._intor)
shls_slice = (0, mol.nbas, 0, mol.nbas, mol.nbas, len(bas))
ngrids = grid_coords.shape[0]
vj = vk = None
fjk = []
dmsptr = []
vjkptr = []
if with_j:
if dms[0].ndim == 1: # the value of density at each grid
vj = numpy.zeros((len(dms),ncomp,nao,nao))[:,0]
for i, dm in enumerate(dms):
dmsptr.append(dm.ctypes.data_as(ctypes.c_void_p))
vjkptr.append(vj[i].ctypes.data_as(ctypes.c_void_p))
fjk.append(_vhf._fpointer('SGXnr'+aosym+'_ijg_g_ij'))
else:
vj = numpy.zeros((len(dms),ncomp,ngrids))[:,0]
for i, dm in enumerate(dms):
dmsptr.append(dm.ctypes.data_as(ctypes.c_void_p))
vjkptr.append(vj[i].ctypes.data_as(ctypes.c_void_p))
fjk.append(_vhf._fpointer('SGXnr'+aosym+'_ijg_ji_g'))
if with_k:
vk = numpy.zeros((len(fg),ncomp,ngrids,nao))[:,0]
for i, dm in enumerate(fg):
dmsptr.append(dm.ctypes.data_as(ctypes.c_void_p))
vjkptr.append(vk[i].ctypes.data_as(ctypes.c_void_p))
fjk.append(_vhf._fpointer('SGXnr'+aosym+'_ijg_gj_gi'))
n_dm = len(fjk)
fjk = (ctypes.c_void_p*(n_dm))(*fjk)
dmsptr = (ctypes.c_void_p*(n_dm))(*dmsptr)
vjkptr = (ctypes.c_void_p*(n_dm))(*vjkptr)
drv(cintor, fdot, fjk, dmsptr, vjkptr, n_dm, ncomp,
(ctypes.c_int*6)(*shls_slice),
ao_loc.ctypes.data_as(ctypes.c_void_p),
sgxopt._cintopt, sgxopt._this,
atm.ctypes.data_as(ctypes.c_void_p), ctypes.c_int(mol.natm),
bas.ctypes.data_as(ctypes.c_void_p), ctypes.c_int(mol.nbas),
env.ctypes.data_as(ctypes.c_void_p))
return vj, vk
def write_gto(mol):
atm = numpy.asarray(mol._atm, dtype=numpy.int32, order='C')
bas = numpy.asarray(mol._bas, dtype=numpy.int32, order='C')
env = numpy.asarray(mol._env, dtype=numpy.double, order='C')
natm = atm.shape[0]
nbas = bas.shape[0]
eval_name, comp = _get_intor_and_comp(mol, "GTOval")
ao_loc = make_loc(bas, eval_name)
shls_slice = (0, nbas)
sh0, sh1 = shls_slice
nao = ao_loc[sh1] - ao_loc[sh0]
ngrids = 1
non0tab = numpy.ones(((ngrids + BLKSIZE - 1) // BLKSIZE, nbas),
dtype=numpy.int8)
basisInfo = open("basisInfo.txt", "w")
if ("cart" in eval_name):
basisInfo.write("cart\n")
else:
basisInfo.write("sph\n")
basisInfo.write("%i\n" % (nao))
basisInfo.write("%i\n%i \n" % (shls_slice[0], shls_slice[1]))
basisInfo.write("%i\n" % (len(ao_loc)))
for a in ao_loc:
basisInfo.write("%i " % (a))
basisInfo.write("\n%i %i\n" % (atm.shape[0], atm.shape[1]))
for i in range(atm.shape[0]):
for j in range(atm.shape[1]):
basisInfo.write("%i " % (atm[i][j]))
basisInfo.write("\n%i %i\n" % (bas.shape[0], bas.shape[1]))
for i in range(bas.shape[0]):
for j in range(bas.shape[1]):
basisInfo.write("%i " % (bas[i][j]))
basisInfo.write("\n%i\n" % (len(env)))
for e in env:
basisInfo.write("%35.18e\n" % (e))
basisInfo.write("%i\n" % (len(non0tab[0])))
for e in non0tab[0]:
basisInfo.write("%i " % (e))
basisInfo.close()
def nr_e1fill(intor, sh_range, atm, bas, env,
aosym='s1', comp=1, ao2mopt=None, out=None):
assert(aosym in ('s4', 's2ij', 's2kl', 's1'))
c_atm = numpy.asarray(atm, dtype=numpy.int32, order='C')
c_bas = numpy.asarray(bas, dtype=numpy.int32, order='C')
c_env = numpy.asarray(env, order='C')
natm = ctypes.c_int(c_atm.shape[0])
nbas = ctypes.c_int(c_bas.shape[0])
ao_loc = moleintor.make_loc(bas, intor)
nao = ao_loc[-1]
klsh0, klsh1, nkl = sh_range
if aosym in ('s4', 's2ij'):
nao_pair = nao * (nao+1) // 2
else:
nao_pair = nao * nao
if out is None:
out = numpy.empty((comp,nkl,nao_pair))
else:
out = numpy.ndarray((comp,nkl,nao_pair), buffer=out)
if out.size == 0:
return out
if ao2mopt is not None:
cao2mopt = ao2mopt._this
cintopt = ao2mopt._cintopt
cintor = ao2mopt._intor
else:
cao2mopt = lib.c_null_ptr()
cintor = _fpointer(intor)
cintopt = _vhf.make_cintopt(c_atm, c_bas, c_env, intor)
fdrv = getattr(libao2mo, 'AO2MOnr_e1fill_drv')
fill = _fpointer('AO2MOfill_nr_' + aosym)
fdrv(cintor, fill, out.ctypes.data_as(ctypes.c_void_p),
ctypes.c_int(klsh0), ctypes.c_int(klsh1-klsh0),
ctypes.c_int(nkl), ctypes.c_int(comp),
ao_loc.ctypes.data_as(ctypes.c_void_p), cintopt, cao2mopt,
c_atm.ctypes.data_as(ctypes.c_void_p), natm,
c_bas.ctypes.data_as(ctypes.c_void_p), nbas,
c_env.ctypes.data_as(ctypes.c_void_p))
return out
def set_dm(self, dm, atm, bas, env):
if self._dmcondname is not None:
c_atm = numpy.asarray(atm, dtype=numpy.int32, order='C')
c_bas = numpy.asarray(bas, dtype=numpy.int32, order='C')
c_env = numpy.asarray(env, dtype=numpy.double, order='C')
natm = ctypes.c_int(c_atm.shape[0])
nbas = ctypes.c_int(c_bas.shape[0])
if isinstance(dm, numpy.ndarray) and dm.ndim == 2:
n_dm = 1
else:
n_dm = len(dm)
dm = numpy.asarray(dm, order='C')
ao_loc = make_loc(c_bas, self._intor)
fsetdm = getattr(libcvhf, self._dmcondname)
fsetdm(self._this, dm.ctypes.data_as(ctypes.c_void_p),
ctypes.c_int(n_dm), ao_loc.ctypes.data_as(ctypes.c_void_p),
c_atm.ctypes.data_as(ctypes.c_void_p), natm,
c_bas.ctypes.data_as(ctypes.c_void_p), nbas,
c_env.ctypes.data_as(ctypes.c_void_p))
def set_dm(self, dm, atm, bas, env):
if self._dmcondname is not None:
c_atm = numpy.asarray(atm, dtype=numpy.int32, order='C')
c_bas = numpy.asarray(bas, dtype=numpy.int32, order='C')
c_env = numpy.asarray(env, dtype=numpy.double, order='C')
natm = ctypes.c_int(c_atm.shape[0])
nbas = ctypes.c_int(c_bas.shape[0])
if isinstance(dm, numpy.ndarray) and dm.ndim == 2:
n_dm = 1
else:
n_dm = len(dm)
dm = numpy.asarray(dm, order='C')
ao_loc = make_loc(c_bas, self._intor)
fsetdm = getattr(libcvhf, self._dmcondname)
fsetdm(self._this,
dm.ctypes.data_as(ctypes.c_void_p), ctypes.c_int(n_dm),
ao_loc.ctypes.data_as(ctypes.c_void_p),
c_atm.ctypes.data_as(ctypes.c_void_p), natm,
c_bas.ctypes.data_as(ctypes.c_void_p), nbas,
c_env.ctypes.data_as(ctypes.c_void_p))
def init_cvhf_direct(self, mol, intor, qcondname):
'''qcondname can be the function pointer or the name of a C function
defined in libcvhf module
'''
intor = mol._add_suffix(intor)
if intor == self._intor:
cintopt = self._cintopt
else:
cintopt = lib.c_null_ptr()
ao_loc = make_loc(mol._bas, intor)
if isinstance(qcondname, ctypes._CFuncPtr):
fsetqcond = qcondname
else:
fsetqcond = getattr(libcvhf, qcondname)
natm = ctypes.c_int(mol.natm)
nbas = ctypes.c_int(mol.nbas)
fsetqcond(self._this, getattr(libcvhf, intor), cintopt,
ao_loc.ctypes.data_as(ctypes.c_void_p),
mol._atm.ctypes.data_as(ctypes.c_void_p), natm,
mol._bas.ctypes.data_as(ctypes.c_void_p), nbas,
mol._env.ctypes.data_as(ctypes.c_void_p))
def init_cvhf_direct(self, mol, intor, prescreen, qcondname):
c_atm = numpy.asarray(mol._atm, dtype=numpy.int32, order='C')
c_bas = numpy.asarray(mol._bas, dtype=numpy.int32, order='C')
c_env = numpy.asarray(mol._env, dtype=numpy.double, order='C')
natm = ctypes.c_int(c_atm.shape[0])
nbas = ctypes.c_int(c_bas.shape[0])
self._cintopt = make_cintopt(c_atm, c_bas, c_env, intor)
libcvhf.CVHFinit_optimizer(ctypes.byref(self._this),
c_atm.ctypes.data_as(ctypes.c_void_p), natm,
c_bas.ctypes.data_as(ctypes.c_void_p), nbas,
c_env.ctypes.data_as(ctypes.c_void_p))
self._this.contents.fprescreen = _fpointer(prescreen)
if prescreen != 'CVHFnoscreen':
ao_loc = make_loc(c_bas, self._intor)
fsetqcond = getattr(libcvhf, qcondname)
fsetqcond(self._this, getattr(libcvhf, intor), self._cintopt,
ao_loc.ctypes.data_as(ctypes.c_void_p),
c_atm.ctypes.data_as(ctypes.c_void_p), natm,
c_bas.ctypes.data_as(ctypes.c_void_p), nbas,
c_env.ctypes.data_as(ctypes.c_void_p))
def init_cvhf_direct(self, mol, intor, prescreen, qcondname):
c_atm = numpy.asarray(mol._atm, dtype=numpy.int32, order='C')
c_bas = numpy.asarray(mol._bas, dtype=numpy.int32, order='C')
c_env = numpy.asarray(mol._env, dtype=numpy.double, order='C')
natm = ctypes.c_int(c_atm.shape[0])
nbas = ctypes.c_int(c_bas.shape[0])
self._cintopt = make_cintopt(c_atm, c_bas, c_env, intor)
libcvhf.CVHFinit_optimizer(ctypes.byref(self._this),
c_atm.ctypes.data_as(ctypes.c_void_p), natm,
c_bas.ctypes.data_as(ctypes.c_void_p), nbas,
c_env.ctypes.data_as(ctypes.c_void_p))
self._this.contents.fprescreen = _fpointer(prescreen)
if prescreen != 'CVHFnoscreen' and qcondname is not None:
ao_loc = make_loc(c_bas, self._intor)
fsetqcond = getattr(libcvhf, qcondname)
fsetqcond(self._this,
getattr(libcvhf, intor), self._cintopt,
ao_loc.ctypes.data_as(ctypes.c_void_p),
c_atm.ctypes.data_as(ctypes.c_void_p), natm,
c_bas.ctypes.data_as(ctypes.c_void_p), nbas,
c_env.ctypes.data_as(ctypes.c_void_p))
def direct_bindm(intor,
aosym,
jkdescript,
dms,
ncomp,
atm,
bas,
env,
vhfopt=None,
cintopt=None,
shls_slice=None):
assert (aosym in ('s8', 's4', 's2ij', 's2kl', 's1', 'aa4', 'a4ij', 'a4kl',
'a2ij', 'a2kl'))
intor = ascint3(intor)
c_atm = numpy.asarray(atm, dtype=numpy.int32, order='C')
c_bas = numpy.asarray(bas, dtype=numpy.int32, order='C')
c_env = numpy.asarray(env, dtype=numpy.double, order='C')
natm = ctypes.c_int(c_atm.shape[0])
nbas = ctypes.c_int(c_bas.shape[0])
if isinstance(dms, numpy.ndarray) and dms.ndim == 2:
dms = dms[numpy.newaxis, :, :]
n_dm = len(dms)
nao = dms[0].shape[0]
dms = [numpy.asarray(dm, order='C') for dm in dms]
if isinstance(jkdescript, str):
jkdescripts = (jkdescript, )
else:
jkdescripts = jkdescript
njk = len(jkdescripts)
assert (njk == n_dm)
if vhfopt is None:
cintor = _fpointer(intor)
cvhfopt = lib.c_null_ptr()
else:
vhfopt.set_dm(dms, atm, bas, env)
cvhfopt = vhfopt._this
cintopt = vhfopt._cintopt
cintor = getattr(libcvhf, vhfopt._intor)
if cintopt is None:
cintopt = make_cintopt(c_atm, c_bas, c_env, intor)
fdrv = getattr(libcvhf, 'CVHFnr_direct_drv')
dotsym = _INTSYMAP[aosym]
fdot = _fpointer('CVHFdot_nr' + dotsym)
if shls_slice is None:
shls_slice = (0, c_bas.shape[0]) * 4
ao_loc = make_loc(bas, intor)
vjk = []
descr_sym = [x.split('->') for x in jkdescripts]
fjk = (ctypes.c_void_p * (n_dm))()
dmsptr = (ctypes.c_void_p * (n_dm))()
vjkptr = (ctypes.c_void_p * (n_dm))()
for i, (dmsym, vsym) in enumerate(descr_sym):
f1 = _fpointer('CVHFnr%s_%s_%s' % (aosym, dmsym, vsym))
if dms[i].shape != get_dims(dmsym, shls_slice, ao_loc):
raise RuntimeError(
'dm[%d] shape %s is inconsistent with the '
'shls_slice shape %s' %
(i, dms[i].shape, get_dims(dmsym, shls_slice, ao_loc)))
vshape = (ncomp, ) + get_dims(vsym[-2:], shls_slice, ao_loc)
vjk.append(numpy.empty(vshape))
dmsptr[i] = dms[i].ctypes.data_as(ctypes.c_void_p)
vjkptr[i] = vjk[i].ctypes.data_as(ctypes.c_void_p)
fjk[i] = f1
shls_slice = (ctypes.c_int * 8)(*shls_slice)
fdrv(cintor, fdot, fjk, dmsptr, vjkptr, ctypes.c_int(n_dm),
ctypes.c_int(ncomp), shls_slice,
ao_loc.ctypes.data_as(ctypes.c_void_p), cintopt, cvhfopt,
c_atm.ctypes.data_as(ctypes.c_void_p), natm,
c_bas.ctypes.data_as(ctypes.c_void_p), nbas,
c_env.ctypes.data_as(ctypes.c_void_p))
if ncomp == 1:
vjk = [v.reshape(v.shape[1:]) for v in vjk]
if isinstance(jkdescript, str):
vjk = vjk[0]
return vjk
def direct_bindm(intor,
aosym,
jkdescript,
dms,
ncomp,
atm,
bas,
env,
vhfopt=None,
shls_slice=None):
assert (aosym in ('s8', 's4', 's2ij', 's2kl', 's1', 'aa4', 'a4ij', 'a4kl',
'a2ij', 'a2kl'))
intor = ascint3(intor)
c_atm = numpy.asarray(atm, dtype=numpy.int32, order='C')
c_bas = numpy.asarray(bas, dtype=numpy.int32, order='C')
c_env = numpy.asarray(env, dtype=numpy.double, order='C')
natm = ctypes.c_int(c_atm.shape[0])
nbas = ctypes.c_int(c_bas.shape[0])
if isinstance(dms, numpy.ndarray) and dms.ndim == 2:
n_dm = 1
nao = dms.shape[0]
dms = (numpy.asarray(dms, order='C'), )
else:
n_dm = len(dms)
nao = dms[0].shape[0]
dms = [numpy.asarray(dm, order='C') for dm in dms]
if isinstance(jkdescript, str):
njk = 1
jkdescript = (jkdescript, )
else:
njk = len(jkdescript)
assert (njk == n_dm)
if vhfopt is None:
cintor = _fpointer(intor)
cintopt = make_cintopt(c_atm, c_bas, c_env, intor)
cvhfopt = pyscf.lib.c_null_ptr()
else:
vhfopt.set_dm(dms, atm, bas, env)
cvhfopt = vhfopt._this
cintopt = vhfopt._cintopt
cintor = getattr(libcvhf, vhfopt._intor)
fdrv = getattr(libcvhf, 'CVHFnr_direct_drv')
dotsym = _INTSYMAP[aosym]
fdot = _fpointer('CVHFdot_nr' + dotsym)
if shls_slice is None:
shls_slice = (0, c_bas.shape[0]) * 4
ao_loc = make_loc(bas, intor)
vjk = []
descr_sym = [x.split('->') for x in jkdescript]
fjk = (ctypes.c_void_p * (n_dm))()
dmsptr = (ctypes.c_void_p * (n_dm))()
vjkptr = (ctypes.c_void_p * (n_dm))()
for i, (dmsym, vsym) in enumerate(descr_sym):
if dmsym in ('ij', 'kl', 'il', 'kj'):
sys.stderr.write(
'not support DM description %s, transpose to %s\n' %
(dmsym, dmsym[::-1]))
dmsym = dmsym[::-1]
f1 = _fpointer('CVHFnr%s_%s_%s' % (aosym, dmsym, vsym))
assert (dms[i].shape == get_dims(dmsym, shls_slice, ao_loc))
vshape = (ncomp, ) + get_dims(vsym[-2:], shls_slice, ao_loc)
vjk.append(numpy.empty(vshape))
dmsptr[i] = dms[i].ctypes.data_as(ctypes.c_void_p)
vjkptr[i] = vjk[i].ctypes.data_as(ctypes.c_void_p)
fjk[i] = f1
shls_slice = (ctypes.c_int * 8)(*shls_slice)
fdrv(cintor, fdot, fjk, dmsptr, vjkptr, ctypes.c_int(n_dm),
ctypes.c_int(ncomp), shls_slice,
ao_loc.ctypes.data_as(ctypes.c_void_p), cintopt, cvhfopt,
c_atm.ctypes.data_as(ctypes.c_void_p), natm,
c_bas.ctypes.data_as(ctypes.c_void_p), nbas,
c_env.ctypes.data_as(ctypes.c_void_p))
if ncomp == 1:
vjk = [v.reshape(v.shape[1:]) for v in vjk]
if njk == 1:
vjk = vjk[0]
return vjk
def direct(dms, atm, bas, env, vhfopt=None, hermi=0, cart=False):
c_atm = numpy.asarray(atm, dtype=numpy.int32, order='C')
c_bas = numpy.asarray(bas, dtype=numpy.int32, order='C')
c_env = numpy.asarray(env, dtype=numpy.double, order='C')
natm = ctypes.c_int(c_atm.shape[0])
nbas = ctypes.c_int(c_bas.shape[0])
if isinstance(dms, numpy.ndarray) and dms.ndim == 2:
n_dm = 1
nao = dms.shape[0]
dms = (numpy.asarray(dms, order='C'), )
else:
n_dm = len(dms)
nao = dms[0].shape[0]
dms = numpy.asarray(dms, order='C')
if vhfopt is None:
if cart:
intor = 'int2e_cart'
else:
intor = 'int2e_sph'
cintopt = make_cintopt(c_atm, c_bas, c_env, 'int2e_sph')
cvhfopt = pyscf.lib.c_null_ptr()
else:
vhfopt.set_dm(dms, atm, bas, env)
cvhfopt = vhfopt._this
cintopt = vhfopt._cintopt
intor = vhfopt._intor
cintor = _fpointer(intor)
fdrv = getattr(libcvhf, 'CVHFnr_direct_drv')
fdot = _fpointer('CVHFdot_nrs8')
fvj = _fpointer('CVHFnrs8_ji_s2kl')
if hermi == 1:
fvk = _fpointer('CVHFnrs8_li_s2kj')
else:
fvk = _fpointer('CVHFnrs8_li_s1kj')
vjk = numpy.empty((2, n_dm, nao, nao))
fjk = (ctypes.c_void_p * (2 * n_dm))()
dmsptr = (ctypes.c_void_p * (2 * n_dm))()
vjkptr = (ctypes.c_void_p * (2 * n_dm))()
for i in range(n_dm):
dmsptr[i] = dms[i].ctypes.data_as(ctypes.c_void_p)
vjkptr[i] = vjk[0, i].ctypes.data_as(ctypes.c_void_p)
fjk[i] = fvj
for i in range(n_dm):
dmsptr[n_dm + i] = dms[i].ctypes.data_as(ctypes.c_void_p)
vjkptr[n_dm + i] = vjk[1, i].ctypes.data_as(ctypes.c_void_p)
fjk[n_dm + i] = fvk
shls_slice = (ctypes.c_int * 8)(*([0, c_bas.shape[0]] * 4))
ao_loc = make_loc(bas, intor)
fdrv(cintor, fdot, fjk, dmsptr, vjkptr, ctypes.c_int(n_dm * 2),
ctypes.c_int(1), shls_slice, ao_loc.ctypes.data_as(ctypes.c_void_p),
cintopt, cvhfopt, c_atm.ctypes.data_as(ctypes.c_void_p), natm,
c_bas.ctypes.data_as(ctypes.c_void_p), nbas,
c_env.ctypes.data_as(ctypes.c_void_p))
# vj must be symmetric
for idm in range(n_dm):
vjk[0, idm] = pyscf.lib.hermi_triu(vjk[0, idm], 1)
if hermi != 0: # vk depends
for idm in range(n_dm):
vjk[1, idm] = pyscf.lib.hermi_triu(vjk[1, idm], hermi)
if n_dm == 1:
vjk = vjk.reshape(2, nao, nao)
return vjk
def intor_cross(intor, cell1, cell2, comp=1, hermi=0, kpts=None, kpt=None):
r'''1-electron integrals from two cells like
.. math::
\langle \mu | intor | \nu \rangle, \mu \in cell1, \nu \in cell2
'''
if kpts is None:
if kpt is not None:
kpts_lst = np.reshape(kpt, (1, 3))
else:
kpts_lst = np.zeros((1, 3))
else:
kpts_lst = np.reshape(kpts, (-1, 3))
nkpts = len(kpts_lst)
atm, bas, env = conc_env(cell1._atm, cell1._bas, cell1._env, cell2._atm,
cell2._bas, cell2._env)
atm = np.asarray(atm, dtype=np.int32)
bas = np.asarray(bas, dtype=np.int32)
env = np.asarray(env, dtype=np.double)
natm = len(atm)
nbas = len(bas)
shls_slice = (0, cell1.nbas, cell1.nbas, nbas)
ao_loc = moleintor.make_loc(bas, intor)
ni = ao_loc[shls_slice[1]] - ao_loc[shls_slice[0]]
nj = ao_loc[shls_slice[3]] - ao_loc[shls_slice[2]]
out = [
np.zeros((ni, nj, comp), order='F', dtype=np.complex128)
for k in range(nkpts)
]
out_ptrs = (ctypes.c_void_p *
nkpts)(*[x.ctypes.data_as(ctypes.c_void_p) for x in out])
if hermi == 0:
aosym = 's1'
else:
aosym = 's2'
if '2c2e' in intor:
fill = getattr(libpbc, 'PBCnr2c2e_fill_' + aosym)
else:
assert ('2e' not in intor)
fill = getattr(libpbc, 'PBCnr2c_fill_' + aosym)
fintor = getattr(moleintor.libcgto, intor)
intopt = lib.c_null_ptr()
Ls = cell1.get_lattice_Ls(rcut=max(cell1.rcut, cell2.rcut))
expLk = np.asarray(np.exp(1j * np.dot(Ls, kpts_lst.T)), order='C')
xyz = np.asarray(cell2.atom_coords(), order='C')
ptr_coords = np.asarray(atm[cell1.natm:, mole.PTR_COORD],
dtype=np.int32,
order='C')
drv = libpbc.PBCnr2c_drv
drv(fintor, fill, out_ptrs, xyz.ctypes.data_as(ctypes.c_void_p),
ptr_coords.ctypes.data_as(ctypes.c_void_p), ctypes.c_int(cell2.natm),
Ls.ctypes.data_as(ctypes.c_void_p), ctypes.c_int(len(Ls)),
expLk.ctypes.data_as(ctypes.c_void_p), ctypes.c_int(nkpts),
ctypes.c_int(comp), (ctypes.c_int * 4)(*(shls_slice[:4])),
ao_loc.ctypes.data_as(ctypes.c_void_p), intopt,
atm.ctypes.data_as(ctypes.c_void_p), ctypes.c_int(natm),
bas.ctypes.data_as(ctypes.c_void_p), ctypes.c_int(nbas),
env.ctypes.data_as(ctypes.c_void_p))
def trans(out):
out = out.transpose(2, 0, 1)
if hermi == lib.HERMITIAN:
# GTOint2c fills the upper triangular of the F-order array.
idx = np.triu_indices(ni)
for i in range(comp):
out[i, idx[1], idx[0]] = out[i, idx[0], idx[1]].conj()
elif hermi == lib.ANTIHERMI:
idx = np.triu_indices(ni)
for i in range(comp):
out[i, idx[1], idx[0]] = -out[i, idx[0], idx[1]].conj()
elif hermi == lib.SYMMETRIC:
idx = np.triu_indices(ni)
for i in range(comp):
out[i, idx[1], idx[0]] = out[i, idx[0], idx[1]]
if comp == 1:
out = out.reshape(ni, nj)
return out
for k, kpt in enumerate(kpts_lst):
if abs(kpt).sum() < 1e-9: # gamma_point
out[k] = np.asarray(trans(out[k].real), order='C')
else:
out[k] = np.asarray(trans(out[k]), order='C')
if kpts is None or np.shape(kpts) == (3, ):
# A single k-point
out = out[0]
return out
def intor_cross(intor, cell1, cell2, comp=1, hermi=0, kpts=None, kpt=None):
r'''1-electron integrals from two cells like
.. math::
\langle \mu | intor | \nu \rangle, \mu \in cell1, \nu \in cell2
'''
if kpts is None:
if kpt is not None:
kpts_lst = np.reshape(kpt, (1,3))
else:
kpts_lst = np.zeros((1,3))
else:
kpts_lst = np.reshape(kpts, (-1,3))
nkpts = len(kpts_lst)
atm, bas, env = conc_env(cell1._atm, cell1._bas, cell1._env,
cell2._atm, cell2._bas, cell2._env)
atm = np.asarray(atm, dtype=np.int32)
bas = np.asarray(bas, dtype=np.int32)
env = np.asarray(env, dtype=np.double)
natm = len(atm)
nbas = len(bas)
shls_slice = (0, cell1.nbas, cell1.nbas, nbas)
ao_loc = moleintor.make_loc(bas, intor)
ni = ao_loc[shls_slice[1]] - ao_loc[shls_slice[0]]
nj = ao_loc[shls_slice[3]] - ao_loc[shls_slice[2]]
out = [np.zeros((ni,nj,comp), order='F', dtype=np.complex128)
for k in range(nkpts)]
out_ptrs = (ctypes.c_void_p*nkpts)(
*[x.ctypes.data_as(ctypes.c_void_p) for x in out])
if hermi == 0:
aosym = 's1'
else:
aosym = 's2'
if '2c2e' in intor:
fill = getattr(libpbc, 'PBCnr2c2e_fill_'+aosym)
else:
assert('2e' not in intor)
fill = getattr(libpbc, 'PBCnr2c_fill_'+aosym)
fintor = moleintor._fpointer(intor)
intopt = lib.c_null_ptr()
nimgs = np.max((cell1.nimgs, cell2.nimgs), axis=0)
Ls = np.asarray(cell1.get_lattice_Ls(nimgs), order='C')
expLk = np.asarray(np.exp(1j*np.dot(Ls, kpts_lst.T)), order='C')
xyz = np.asarray(cell2.atom_coords(), order='C')
ptr_coords = np.asarray(atm[cell1.natm:,mole.PTR_COORD],
dtype=np.int32, order='C')
drv = libpbc.PBCnr2c_drv
drv(fintor, fill, out_ptrs, xyz.ctypes.data_as(ctypes.c_void_p),
ptr_coords.ctypes.data_as(ctypes.c_void_p), ctypes.c_int(cell2.natm),
Ls.ctypes.data_as(ctypes.c_void_p), ctypes.c_int(len(Ls)),
expLk.ctypes.data_as(ctypes.c_void_p), ctypes.c_int(nkpts),
ctypes.c_int(comp), (ctypes.c_int*4)(*(shls_slice[:4])),
ao_loc.ctypes.data_as(ctypes.c_void_p), intopt,
atm.ctypes.data_as(ctypes.c_void_p), ctypes.c_int(natm),
bas.ctypes.data_as(ctypes.c_void_p), ctypes.c_int(nbas),
env.ctypes.data_as(ctypes.c_void_p))
def trans(out):
out = out.transpose(2,0,1)
if hermi == lib.HERMITIAN:
# GTOint2c fills the upper triangular of the F-order array.
idx = np.triu_indices(ni)
for i in range(comp):
out[i,idx[1],idx[0]] = out[i,idx[0],idx[1]].conj()
elif hermi == lib.ANTIHERMI:
idx = np.triu_indices(ni)
for i in range(comp):
out[i,idx[1],idx[0]] = -out[i,idx[0],idx[1]].conj()
elif hermi == lib.SYMMETRIC:
idx = np.triu_indices(ni)
for i in range(comp):
out[i,idx[1],idx[0]] = out[i,idx[0],idx[1]]
if comp == 1:
out = out.reshape(ni,nj)
return out
for k, kpt in enumerate(kpts_lst):
if abs(kpt).sum() < 1e-9: # gamma_point
out[k] = np.asarray(trans(out[k].real), order='C')
else:
out[k] = np.asarray(trans(out[k]), order='C')
if kpts is None or np.shape(kpts) == (3,):
# A single k-point
out = out[0]
return out
def intor_cross(intor, cell1, cell2, comp=1, hermi=0, kpts=None, kpt=None):
r'''1-electron integrals from two cells like
.. math::
\langle \mu | intor | \nu \rangle, \mu \in cell1, \nu \in cell2
'''
intor = moleintor.ascint3(intor)
if kpts is None:
if kpt is not None:
kpts_lst = np.reshape(kpt, (1, 3))
else:
kpts_lst = np.zeros((1, 3))
else:
kpts_lst = np.reshape(kpts, (-1, 3))
nkpts = len(kpts_lst)
atm, bas, env = conc_env(cell1._atm, cell1._bas, cell1._env, cell2._atm,
cell2._bas, cell2._env)
atm = np.asarray(atm, dtype=np.int32)
bas = np.asarray(bas, dtype=np.int32)
env = np.asarray(env, dtype=np.double)
natm = len(atm)
nbas = len(bas)
shls_slice = (0, cell1.nbas, cell1.nbas, nbas)
ao_loc = moleintor.make_loc(bas, intor)
ni = ao_loc[shls_slice[1]] - ao_loc[shls_slice[0]]
nj = ao_loc[shls_slice[3]] - ao_loc[shls_slice[2]]
out = np.empty((nkpts, comp, ni, nj), dtype=np.complex128)
if hermi == 0:
aosym = 's1'
else:
aosym = 's2'
fill = getattr(libpbc, 'PBCnr2c_fill_k' + aosym)
fintor = getattr(moleintor.libcgto, intor)
intopt = lib.c_null_ptr()
Ls = cell1.get_lattice_Ls(rcut=max(cell1.rcut, cell2.rcut))
expkL = np.asarray(np.exp(1j * np.dot(kpts_lst, Ls.T)), order='C')
drv = libpbc.PBCnr2c_drv
drv(fintor, fill, out.ctypes.data_as(ctypes.c_void_p), ctypes.c_int(nkpts),
ctypes.c_int(comp), ctypes.c_int(len(Ls)),
Ls.ctypes.data_as(ctypes.c_void_p),
expkL.ctypes.data_as(ctypes.c_void_p),
(ctypes.c_int * 4)(*(shls_slice[:4])),
ao_loc.ctypes.data_as(ctypes.c_void_p), intopt,
atm.ctypes.data_as(ctypes.c_void_p), ctypes.c_int(natm),
bas.ctypes.data_as(ctypes.c_void_p), ctypes.c_int(nbas),
env.ctypes.data_as(ctypes.c_void_p))
mat = []
for k, kpt in enumerate(kpts_lst):
v = out[k]
if hermi != 0:
for ic in range(comp):
lib.hermi_triu(v[ic], hermi=hermi, inplace=True)
if comp == 1:
v = v[0]
if abs(kpt).sum() < 1e-9: # gamma_point
v = v.real
mat.append(v)
if kpts is None or np.shape(kpts) == (3, ): # A single k-point
mat = mat[0]
return mat
def eval_gto(mol,
eval_name,
coords,
comp=None,
shls_slice=None,
non0tab=None,
ao_loc=None,
out=None):
r'''Evaluate AO function value on the given grids,
Args:
eval_name : str
================== ====== =======================
Function comp Expression
================== ====== =======================
"GTOval_sph" 1 |AO>
"GTOval_ip_sph" 3 nabla |AO>
"GTOval_ig_sph" 3 (#C(0 1) g) |AO>
"GTOval_ipig_sph" 3 (#C(0 1) nabla g) |AO>
"GTOval_cart" 1 |AO>
"GTOval_ip_cart" 3 nabla |AO>
"GTOval_ig_cart" 3 (#C(0 1) g)|AO>
================== ====== =======================
atm : int32 ndarray
libcint integral function argument
bas : int32 ndarray
libcint integral function argument
env : float64 ndarray
libcint integral function argument
coords : 2D array, shape (N,3)
The coordinates of the grids.
Kwargs:
comp : int
Number of the components of the operator
shls_slice : 2-element list
(shl_start, shl_end).
If given, only part of AOs (shl_start <= shell_id < shl_end) are
evaluated. By default, all shells defined in mol will be evaluated.
non0tab : 2D bool array
mask array to indicate whether the AO values are zero. The mask
array can be obtained by calling :func:`dft.gen_grid.make_mask`
out : ndarray
If provided, results are written into this array.
Returns:
2D array of shape (N,nao) Or 3D array of shape (\*,N,nao) to store AO
values on grids.
Examples:
>>> mol = gto.M(atom='O 0 0 0; H 0 0 1; H 0 1 0', basis='ccpvdz')
>>> coords = numpy.random.random((100,3)) # 100 random points
>>> ao_value = mol.eval_gto("GTOval_sph", coords)
>>> print(ao_value.shape)
(100, 24)
>>> ao_value = mol.eval_gto("GTOval_ig_sph", coords)
>>> print(ao_value.shape)
(3, 100, 24)
'''
eval_name, comp = _get_intor_and_comp(mol, eval_name, comp)
atm = numpy.asarray(mol._atm, dtype=numpy.int32, order='C')
bas = numpy.asarray(mol._bas, dtype=numpy.int32, order='C')
env = numpy.asarray(mol._env, dtype=numpy.double, order='C')
coords = numpy.asarray(coords, dtype=numpy.double, order='F')
natm = atm.shape[0]
nbas = bas.shape[0]
ngrids = coords.shape[0]
if ao_loc is None:
ao_loc = make_loc(bas, eval_name)
if shls_slice is None:
shls_slice = (0, nbas)
sh0, sh1 = shls_slice
nao = ao_loc[sh1] - ao_loc[sh0]
if 'spinor' in eval_name:
ao = numpy.ndarray((2, comp, nao, ngrids),
dtype=numpy.complex128,
buffer=out)
else:
ao = numpy.ndarray((comp, nao, ngrids), buffer=out)
if non0tab is None:
non0tab = numpy.ones(((ngrids + BLKSIZE - 1) // BLKSIZE, nbas),
dtype=numpy.int8)
drv = getattr(libcgto, eval_name)
drv(ctypes.c_int(ngrids), (ctypes.c_int * 2)(*shls_slice),
ao_loc.ctypes.data_as(ctypes.c_void_p),
ao.ctypes.data_as(ctypes.c_void_p),
coords.ctypes.data_as(ctypes.c_void_p),
non0tab.ctypes.data_as(ctypes.c_void_p),
atm.ctypes.data_as(ctypes.c_void_p), ctypes.c_int(natm),
bas.ctypes.data_as(ctypes.c_void_p), ctypes.c_int(nbas),
env.ctypes.data_as(ctypes.c_void_p))
ao = numpy.swapaxes(ao, -1, -2)
if comp == 1:
if 'spinor' in eval_name:
ao = ao[:, 0]
else:
ao = ao[0]
return ao
def eval_gto(cell,
eval_name,
coords,
comp=None,
kpts=None,
kpt=None,
shls_slice=None,
non0tab=None,
ao_loc=None,
out=None,
Ls=None,
rcut=None):
r'''Evaluate PBC-AO function value on the given grids,
Args:
eval_name : str
========================== =======================
Function Expression
========================== =======================
"GTOval_sph" \sum_T exp(ik*T) |AO>
"GTOval_ip_sph" nabla \sum_T exp(ik*T) |AO>
"GTOval_cart" \sum_T exp(ik*T) |AO>
"GTOval_ip_cart" nabla \sum_T exp(ik*T) |AO>
========================== =======================
atm : int32 ndarray
libcint integral function argument
bas : int32 ndarray
libcint integral function argument
env : float64 ndarray
libcint integral function argument
coords : 2D array, shape (N,3)
The coordinates of the grids.
Kwargs:
shls_slice : 2-element list
(shl_start, shl_end).
If given, only part of AOs (shl_start <= shell_id < shl_end) are
evaluated. By default, all shells defined in cell will be evaluated.
non0tab : 2D bool array
mask array to indicate whether the AO values are zero. The mask
array can be obtained by calling :func:`dft.gen_grid.make_mask`
out : ndarray
If provided, results are written into this array.
Returns:
A list of 2D (or 3D) arrays to hold the AO values on grids. Each
element of the list corresponds to a k-point and it has the shape
(N,nao) Or shape (\*,N,nao).
Examples:
>>> cell = pbc.gto.M(a=numpy.eye(3)*4, atom='He 1 1 1', basis='6-31g')
>>> coords = cell.get_uniform_grids([20,20,20])
>>> kpts = cell.make_kpts([3,3,3])
>>> ao_value = cell.pbc_eval_gto("GTOval_sph", coords, kpts)
>>> len(ao_value)
27
>>> ao_value[0].shape
(100, 2)
>>> ao_value = cell.pbc_eval_gto("GTOval_ig_sph", coords, kpts, comp=3)
>>> print(ao_value.shape)
>>> len(ao_value)
27
>>> ao_value[0].shape
(3, 100, 2)
'''
if eval_name[:3] == 'PBC': # PBCGTOval_xxx
eval_name, comp = _get_intor_and_comp(cell, eval_name[3:], comp)
else:
eval_name, comp = _get_intor_and_comp(cell, eval_name, comp)
eval_name = 'PBC' + eval_name
atm = numpy.asarray(cell._atm, dtype=numpy.int32, order='C')
bas = numpy.asarray(cell._bas, dtype=numpy.int32, order='C')
env = numpy.asarray(cell._env, dtype=numpy.double, order='C')
natm = atm.shape[0]
nbas = bas.shape[0]
if kpts is None:
if kpt is not None:
kpts_lst = numpy.reshape(kpt, (1, 3))
else:
kpts_lst = numpy.zeros((1, 3))
else:
kpts_lst = numpy.reshape(kpts, (-1, 3))
nkpts = len(kpts_lst)
ngrids = len(coords)
if non0tab is None:
non0tab = numpy.empty(((ngrids + BLKSIZE - 1) // BLKSIZE, nbas),
dtype=numpy.uint8)
# non0tab stores the number of images to be summed in real space.
# Initializing it to 255 means all images should be included
non0tab[:] = 0xff
if ao_loc is None:
ao_loc = moleintor.make_loc(bas, eval_name)
if shls_slice is None:
shls_slice = (0, nbas)
sh0, sh1 = shls_slice
nao = ao_loc[sh1] - ao_loc[sh0]
out = numpy.empty((nkpts, comp, nao, ngrids), dtype=numpy.complex128)
coords = numpy.asarray(coords, order='F')
# For atoms near the boundary of the cell, it is necessary (even in low-
# dimensional systems) to include lattice translations in all 3 dimensions.
if Ls is None:
if cell.dimension < 2 or cell.low_dim_ft_type == 'inf_vacuum':
Ls = cell.get_lattice_Ls(dimension=cell.dimension)
else:
Ls = cell.get_lattice_Ls(dimension=3)
Ls = Ls[numpy.argsort(lib.norm(Ls, axis=1))]
expLk = numpy.exp(1j * numpy.asarray(numpy.dot(Ls, kpts_lst.T), order='C'))
if rcut is None:
rcut = _estimate_rcut(cell)
drv = getattr(libpbc, eval_name)
drv(ctypes.c_int(ngrids), (ctypes.c_int * 2)(*shls_slice),
ao_loc.ctypes.data_as(ctypes.c_void_p),
Ls.ctypes.data_as(ctypes.c_void_p), ctypes.c_int(len(Ls)),
expLk.ctypes.data_as(ctypes.c_void_p), ctypes.c_int(nkpts),
out.ctypes.data_as(ctypes.c_void_p),
coords.ctypes.data_as(ctypes.c_void_p),
rcut.ctypes.data_as(ctypes.c_void_p),
non0tab.ctypes.data_as(ctypes.c_void_p),
atm.ctypes.data_as(ctypes.c_void_p), ctypes.c_int(natm),
bas.ctypes.data_as(ctypes.c_void_p), ctypes.c_int(nbas),
env.ctypes.data_as(ctypes.c_void_p))
ao_kpts = []
for k, kpt in enumerate(kpts_lst):
v = out[k]
if abs(kpt).sum() < 1e-9:
v = numpy.asarray(v.real, order='C')
v = v.transpose(0, 2, 1)
if comp == 1:
v = v[0]
ao_kpts.append(v)
if kpts is None or numpy.shape(kpts) == (3, ): # A single k-point
ao_kpts = ao_kpts[0]
return ao_kpts
def rdirect_bindm(intor,
aosym,
jkdescript,
dms,
ncomp,
atm,
bas,
env,
vhfopt=None,
shls_slice=None):
assert (aosym in ('s8', 's4', 's2ij', 's2kl', 's1', 'a4ij', 'a4kl', 'a2ij',
'a2kl'))
intor = ascint3(intor)
c_atm = numpy.asarray(atm, dtype=numpy.int32, order='C')
c_bas = numpy.asarray(bas, dtype=numpy.int32, order='C')
c_env = numpy.asarray(env, dtype=numpy.double, order='C')
natm = ctypes.c_int(c_atm.shape[0])
nbas = ctypes.c_int(c_bas.shape[0])
if isinstance(dms, numpy.ndarray) and dms.ndim == 2:
n_dm = 1
nao = dms.shape[0]
dms = (numpy.asarray(dms, order='C', dtype=numpy.complex128), )
else:
n_dm = len(dms)
nao = dms[0].shape[0]
dms = numpy.asarray(dms, order='C', dtype=numpy.complex128)
if isinstance(jkdescript, str):
njk = 1
jkdescript = (jkdescript, )
else:
njk = len(jkdescript)
assert (njk == n_dm)
if vhfopt is None:
cintor = _fpointer(intor)
cintopt = make_cintopt(c_atm, c_bas, c_env, intor)
cvhfopt = pyscf.lib.c_null_ptr()
else:
vhfopt.set_dm(dms, atm, bas, env)
cvhfopt = vhfopt._this
cintopt = vhfopt._cintopt
cintor = getattr(libcvhf, vhfopt._intor)
fdrv = getattr(libcvhf, 'CVHFr_direct_drv')
dotsym = _INTSYMAP[aosym]
fdot = _fpointer('CVHFdot_r' + dotsym)
if shls_slice is None:
shls_slice = (0, c_bas.shape[0]) * 4
else:
raise NotImplementedError
ao_loc = make_loc(bas, intor)
unpackas = _INTUNPACKMAP_R[aosym]
descr_sym = [x.split('->') for x in jkdescript]
fjk = (ctypes.c_void_p * (n_dm))()
dm1 = (ctypes.c_void_p * (n_dm))()
for i, (dmsym, vsym) in enumerate(descr_sym):
f1 = _fpointer('CVHFr%s_%s_%s' % (unpackas, dmsym, vsym))
dm1[i] = dms[i].ctypes.data_as(ctypes.c_void_p)
fjk[i] = f1
vjk = numpy.empty((njk, ncomp, nao, nao), dtype=numpy.complex)
fdrv(cintor, fdot, fjk, dm1, vjk.ctypes.data_as(ctypes.c_void_p),
ctypes.c_int(n_dm),
ctypes.c_int(ncomp), (ctypes.c_int * 8)(*shls_slice),
ao_loc.ctypes.data_as(ctypes.c_void_p), cintopt, cvhfopt,
c_atm.ctypes.data_as(ctypes.c_void_p), natm,
c_bas.ctypes.data_as(ctypes.c_void_p), nbas,
c_env.ctypes.data_as(ctypes.c_void_p))
if ncomp == 1:
vjk = vjk.reshape(njk, nao, nao)
if njk == 1:
vjk = vjk.reshape(vjk.shape[1:])
return vjk
def direct(dms,
atm,
bas,
env,
vhfopt=None,
hermi=0,
cart=False,
with_j=True,
with_k=True):
c_atm = numpy.asarray(atm, dtype=numpy.int32, order='C')
c_bas = numpy.asarray(bas, dtype=numpy.int32, order='C')
c_env = numpy.asarray(env, dtype=numpy.double, order='C')
natm = ctypes.c_int(c_atm.shape[0])
nbas = ctypes.c_int(c_bas.shape[0])
dms = numpy.asarray(dms, order='C')
dms_shape = dms.shape
nao = dms_shape[-1]
dms = dms.reshape(-1, nao, nao)
n_dm = dms.shape[0]
if vhfopt is None:
if cart:
intor = 'int2e_cart'
else:
intor = 'int2e_sph'
cintopt = make_cintopt(c_atm, c_bas, c_env, intor)
cvhfopt = lib.c_null_ptr()
else:
vhfopt.set_dm(dms, atm, bas, env)
cvhfopt = vhfopt._this
cintopt = vhfopt._cintopt
intor = vhfopt._intor
cintor = _fpointer(intor)
fdrv = getattr(libcvhf, 'CVHFnr_direct_drv')
fdot = _fpointer('CVHFdot_nrs8')
vj = vk = None
dmsptr = []
vjkptr = []
fjk = []
if with_j:
fvj = _fpointer('CVHFnrs8_ji_s2kl')
vj = numpy.empty((n_dm, nao, nao))
for i, dm in enumerate(dms):
dmsptr.append(dm.ctypes.data_as(ctypes.c_void_p))
vjkptr.append(vj[i].ctypes.data_as(ctypes.c_void_p))
fjk.append(fvj)
if with_k:
if hermi == 1:
fvk = _fpointer('CVHFnrs8_li_s2kj')
else:
fvk = _fpointer('CVHFnrs8_li_s1kj')
vk = numpy.empty((n_dm, nao, nao))
for i, dm in enumerate(dms):
dmsptr.append(dm.ctypes.data_as(ctypes.c_void_p))
vjkptr.append(vk[i].ctypes.data_as(ctypes.c_void_p))
fjk.append(fvk)
shls_slice = (ctypes.c_int * 8)(*([0, c_bas.shape[0]] * 4))
ao_loc = make_loc(bas, intor)
n_ops = len(dmsptr)
comp = 1
fdrv(cintor, fdot, (ctypes.c_void_p * n_ops)(*fjk),
(ctypes.c_void_p * n_ops)(*dmsptr),
(ctypes.c_void_p * n_ops)(*vjkptr), ctypes.c_int(n_ops),
ctypes.c_int(comp), shls_slice,
ao_loc.ctypes.data_as(ctypes.c_void_p), cintopt, cvhfopt,
c_atm.ctypes.data_as(ctypes.c_void_p), natm,
c_bas.ctypes.data_as(ctypes.c_void_p), nbas,
c_env.ctypes.data_as(ctypes.c_void_p))
if with_j:
# vj must be symmetric
for i in range(n_dm):
lib.hermi_triu(vj[i], 1, inplace=True)
vj = vj.reshape(dms_shape)
if with_k:
if hermi != 0:
for i in range(n_dm):
lib.hermi_triu(vk[i], hermi, inplace=True)
vk = vk.reshape(dms_shape)
return vj, vk
def direct(dms, atm, bas, env, vhfopt=None, hermi=0, cart=False):
c_atm = numpy.asarray(atm, dtype=numpy.int32, order='C')
c_bas = numpy.asarray(bas, dtype=numpy.int32, order='C')
c_env = numpy.asarray(env, dtype=numpy.double, order='C')
natm = ctypes.c_int(c_atm.shape[0])
nbas = ctypes.c_int(c_bas.shape[0])
if isinstance(dms, numpy.ndarray) and dms.ndim == 2:
dms = dms[numpy.newaxis,:,:]
n_dm = len(dms)
nao = dms[0].shape[0]
dms = numpy.asarray(dms, order='C')
if vhfopt is None:
if cart:
intor = 'int2e_cart'
else:
intor = 'int2e_sph'
cintopt = make_cintopt(c_atm, c_bas, c_env, intor)
cvhfopt = lib.c_null_ptr()
else:
vhfopt.set_dm(dms, atm, bas, env)
cvhfopt = vhfopt._this
cintopt = vhfopt._cintopt
intor = vhfopt._intor
cintor = _fpointer(intor)
fdrv = getattr(libcvhf, 'CVHFnr_direct_drv')
fdot = _fpointer('CVHFdot_nrs8')
fvj = _fpointer('CVHFnrs8_ji_s2kl')
if hermi == 1:
fvk = _fpointer('CVHFnrs8_li_s2kj')
else:
fvk = _fpointer('CVHFnrs8_li_s1kj')
vjk = numpy.empty((2,n_dm,nao,nao))
fjk = (ctypes.c_void_p*(2*n_dm))()
dmsptr = (ctypes.c_void_p*(2*n_dm))()
vjkptr = (ctypes.c_void_p*(2*n_dm))()
for i in range(n_dm):
dmsptr[i] = dms[i].ctypes.data_as(ctypes.c_void_p)
vjkptr[i] = vjk[0,i].ctypes.data_as(ctypes.c_void_p)
fjk[i] = fvj
for i in range(n_dm):
dmsptr[n_dm+i] = dms[i].ctypes.data_as(ctypes.c_void_p)
vjkptr[n_dm+i] = vjk[1,i].ctypes.data_as(ctypes.c_void_p)
fjk[n_dm+i] = fvk
shls_slice = (ctypes.c_int*8)(*([0, c_bas.shape[0]]*4))
ao_loc = make_loc(bas, intor)
fdrv(cintor, fdot, fjk, dmsptr, vjkptr,
ctypes.c_int(n_dm*2), ctypes.c_int(1),
shls_slice, ao_loc.ctypes.data_as(ctypes.c_void_p), cintopt, cvhfopt,
c_atm.ctypes.data_as(ctypes.c_void_p), natm,
c_bas.ctypes.data_as(ctypes.c_void_p), nbas,
c_env.ctypes.data_as(ctypes.c_void_p))
# vj must be symmetric
for idm in range(n_dm):
vjk[0,idm] = lib.hermi_triu(vjk[0,idm], 1)
if hermi != 0: # vk depends
for idm in range(n_dm):
vjk[1,idm] = lib.hermi_triu(vjk[1,idm], hermi)
if n_dm == 1:
vjk = vjk.reshape(2,nao,nao)
return vjk
def rdirect_bindm(intor, aosym, jkdescript,
dms, ncomp, atm, bas, env, vhfopt=None, cintopt=None,
shls_slice=None):
assert(aosym in ('s8', 's4', 's2ij', 's2kl', 's1',
'a4ij', 'a4kl', 'a2ij', 'a2kl'))
intor = ascint3(intor)
c_atm = numpy.asarray(atm, dtype=numpy.int32, order='C')
c_bas = numpy.asarray(bas, dtype=numpy.int32, order='C')
c_env = numpy.asarray(env, dtype=numpy.double, order='C')
natm = ctypes.c_int(c_atm.shape[0])
nbas = ctypes.c_int(c_bas.shape[0])
if isinstance(dms, numpy.ndarray) and dms.ndim == 2:
dms = dms[numpy.newaxis,:,:]
n_dm = len(dms)
nao = dms[0].shape[0]
dms = numpy.asarray(dms, order='C', dtype=numpy.complex128)
if isinstance(jkdescript, str):
jkdescript = (jkdescript,)
njk = len(jkdescript)
assert(njk == n_dm)
if vhfopt is None:
cintor = _fpointer(intor)
cvhfopt = lib.c_null_ptr()
else:
vhfopt.set_dm(dms, atm, bas, env)
cvhfopt = vhfopt._this
cintopt = vhfopt._cintopt
cintor = getattr(libcvhf, vhfopt._intor)
if cintopt is None:
cintopt = make_cintopt(c_atm, c_bas, c_env, intor)
fdrv = getattr(libcvhf, 'CVHFr_direct_drv')
dotsym = _INTSYMAP[aosym]
fdot = _fpointer('CVHFdot_r'+dotsym)
if shls_slice is None:
shls_slice = (0, c_bas.shape[0])*4
else:
raise NotImplementedError
ao_loc = make_loc(bas, intor)
unpackas = _INTUNPACKMAP_R[aosym]
descr_sym = [x.split('->') for x in jkdescript]
fjk = (ctypes.c_void_p*(n_dm))()
dm1 = (ctypes.c_void_p*(n_dm))()
for i, (dmsym, vsym) in enumerate(descr_sym):
f1 = _fpointer('CVHFr%s_%s_%s'%(unpackas, dmsym, vsym))
dm1[i] = dms[i].ctypes.data_as(ctypes.c_void_p)
fjk[i] = f1
vjk = numpy.empty((njk,ncomp,nao,nao), dtype=numpy.complex)
fdrv(cintor, fdot, fjk, dm1,
vjk.ctypes.data_as(ctypes.c_void_p),
ctypes.c_int(n_dm), ctypes.c_int(ncomp),
(ctypes.c_int*8)(*shls_slice),
ao_loc.ctypes.data_as(ctypes.c_void_p), cintopt, cvhfopt,
c_atm.ctypes.data_as(ctypes.c_void_p), natm,
c_bas.ctypes.data_as(ctypes.c_void_p), nbas,
c_env.ctypes.data_as(ctypes.c_void_p))
if ncomp == 1:
vjk = vjk.reshape(njk,nao,nao)
if njk == 1:
vjk = vjk.reshape(vjk.shape[1:])
return vjk
def get_jk_favorj(sgx,
dm,
hermi=1,
with_j=True,
with_k=True,
direct_scf_tol=1e-13):
t0 = time.clock(), time.time()
mol = sgx.mol
grids = sgx.grids
gthrd = sgx.grids_thrd
dms = numpy.asarray(dm)
dm_shape = dms.shape
nao = dm_shape[-1]
dms = dms.reshape(-1, nao, nao)
nset = dms.shape[0]
if sgx.debug:
batch_nuc = _gen_batch_nuc(mol)
else:
batch_jk = _gen_jk_direct(mol, 's2', with_j, with_k, direct_scf_tol)
# for basis set to shell
intor = mol._add_suffix('int3c2e')
fakemol = gto.fakemol_for_charges(grids.coords)
atm, bas, env = gto.mole.conc_env(mol._atm, mol._bas, mol._env,
fakemol._atm, fakemol._bas, fakemol._env)
ao_loc = moleintor.make_loc(bas, intor)
rao_loc = numpy.zeros((nao), dtype=int)
for i in range(mol.nbas):
for j in range(ao_loc[i], ao_loc[i + 1]):
rao_loc[j] = i
sn = numpy.zeros((nao, nao))
ngrids = grids.coords.shape[0]
max_memory = sgx.max_memory - lib.current_memory()[0]
sblk = sgx.blockdim
blksize = min(ngrids, max(4, int(min(sblk,
max_memory * 1e6 / 8 / nao**2))))
for i0, i1 in lib.prange(0, ngrids, blksize):
coords = grids.coords[i0:i1]
ao = mol.eval_gto('GTOval', coords)
wao = ao * grids.weights[i0:i1, None]
sn += lib.dot(ao.T, wao)
ovlp = mol.intor_symmetric('int1e_ovlp')
proj = scipy.linalg.solve(sn, ovlp)
proj_dm = lib.einsum('ki,xij->xkj', proj, dms)
t1 = logger.timer_debug1(mol, "sgX initialziation", *t0)
vj = numpy.zeros_like(dms)
vk = numpy.zeros_like(dms)
tnuc = 0, 0
for i0, i1 in lib.prange(0, ngrids, blksize):
coords = grids.coords[i0:i1]
ao = mol.eval_gto('GTOval', coords)
wao = ao * grids.weights[i0:i1, None]
fg = lib.einsum('gi,xij->xgj', wao, proj_dm)
mask = numpy.zeros(i1 - i0, dtype=bool)
for i in range(nset):
gmaxfg = numpy.amax(numpy.absolute(fg[i]), axis=1)
gmaxwao_v = numpy.amax(numpy.absolute(ao), axis=1)
gmaxtt = gmaxfg * gmaxwao_v
mask |= numpy.any(gmaxtt > 1e-7)
mask |= numpy.any(gmaxtt < -1e-7)
if not numpy.all(mask):
ao = ao[mask]
wao = wao[mask]
fg = fg[:, mask]
coords = coords[mask]
# screening u by value of grids
umaxg = numpy.amax(numpy.absolute(wao), axis=0)
usi = numpy.argwhere(umaxg > 1e-7).reshape(-1)
if len(usi) != 0:
# screening v by ovlp
uovl = ovlp[usi, :]
vmaxu = numpy.amax(numpy.absolute(uovl), axis=0)
osi = numpy.argwhere(vmaxu > 1e-4).reshape(-1)
udms = proj_dm[0][usi, :]
# screening v by dm and ovlp then triangle matrix bn
dmaxg = numpy.amax(numpy.absolute(udms), axis=0)
dsi = numpy.argwhere(dmaxg > 1e-4).reshape(-1)
vsi = numpy.intersect1d(dsi, osi)
if len(vsi) != 0:
vsh = numpy.unique(rao_loc[vsi])
mol._bvv = vsh
# screening u by value of grids
umaxg = numpy.amax(numpy.absolute(wao), axis=0)
usi = numpy.argwhere(umaxg > 1e-7).reshape(-1)
if len(usi) != 0:
# screening v by ovlp
uovl = ovlp[usi, :]
vmaxu = numpy.amax(numpy.absolute(uovl), axis=0)
osi = numpy.argwhere(vmaxu > 1e-4).reshape(-1)
if len(osi) != 0:
vsh = numpy.unique(rao_loc[osi])
#print(vsh.shape,'eew',vsh)
mol._bvv = vsh
fg = lib.einsum('gi,xij->xgj', wao, proj_dm)
mask = numpy.zeros(i1 - i0, dtype=bool)
for i in range(nset):
mask |= numpy.any(fg[i] > gthrd, axis=1)
mask |= numpy.any(fg[i] < -gthrd, axis=1)
if not numpy.all(mask):
ao = ao[mask]
fg = fg[:, mask]
coords = coords[mask]
if with_j:
rhog = numpy.einsum('xgu,gu->xg', fg, ao)
else:
rhog = None
if sgx.debug:
tnuc = tnuc[0] - time.clock(), tnuc[1] - time.time()
gbn = batch_nuc(mol, coords)
tnuc = tnuc[0] + time.clock(), tnuc[1] + time.time()
if with_j:
jpart = numpy.einsum('guv,xg->xuv', gbn, rhog)
if with_k:
gv = lib.einsum('gtv,xgt->xgv', gbn, fg)
gbn = None
else:
tnuc = tnuc[0] - time.clock(), tnuc[1] - time.time()
jpart, gv = batch_jk(mol, coords, rhog, fg)
tnuc = tnuc[0] + time.clock(), tnuc[1] + time.time()
if with_j:
vj += jpart
if with_k:
for i in range(nset):
vk[i] += lib.einsum('gu,gv->uv', ao, gv[i])
jpart = gv = None
t2 = logger.timer_debug1(mol, "sgX J/K builder", *t1)
tdot = t2[0] - t1[0] - tnuc[0], t2[1] - t1[1] - tnuc[1]
logger.debug1(
sgx, '(CPU, wall) time for integrals (%.2f, %.2f); '
'for tensor contraction (%.2f, %.2f)', tnuc[0], tnuc[1], tdot[0],
tdot[1])
for i in range(nset):
lib.hermi_triu(vj[i], inplace=True)
if with_k and hermi == 1:
vk = (vk + vk.transpose(0, 2, 1)) * .5
logger.timer(mol, "vj and vk", *t0)
return vj.reshape(dm_shape), vk.reshape(dm_shape)
def direct_bindm(intor, aosym, jkdescript,
dms, ncomp, atm, bas, env, vhfopt=None, cintopt=None,
shls_slice=None):
assert(aosym in ('s8', 's4', 's2ij', 's2kl', 's1',
'aa4', 'a4ij', 'a4kl', 'a2ij', 'a2kl'))
intor = ascint3(intor)
c_atm = numpy.asarray(atm, dtype=numpy.int32, order='C')
c_bas = numpy.asarray(bas, dtype=numpy.int32, order='C')
c_env = numpy.asarray(env, dtype=numpy.double, order='C')
natm = ctypes.c_int(c_atm.shape[0])
nbas = ctypes.c_int(c_bas.shape[0])
if isinstance(dms, numpy.ndarray) and dms.ndim == 2:
dms = dms[numpy.newaxis,:,:]
n_dm = len(dms)
nao = dms[0].shape[0]
dms = [numpy.asarray(dm, order='C') for dm in dms]
if isinstance(jkdescript, str):
jkdescripts = (jkdescript,)
else:
jkdescripts = jkdescript
njk = len(jkdescripts)
assert(njk == n_dm)
if vhfopt is None:
cintor = _fpointer(intor)
cvhfopt = lib.c_null_ptr()
else:
vhfopt.set_dm(dms, atm, bas, env)
cvhfopt = vhfopt._this
cintopt = vhfopt._cintopt
cintor = getattr(libcvhf, vhfopt._intor)
if cintopt is None:
cintopt = make_cintopt(c_atm, c_bas, c_env, intor)
fdrv = getattr(libcvhf, 'CVHFnr_direct_drv')
dotsym = _INTSYMAP[aosym]
fdot = _fpointer('CVHFdot_nr'+dotsym)
if shls_slice is None:
shls_slice = (0, c_bas.shape[0])*4
ao_loc = make_loc(bas, intor)
vjk = []
descr_sym = [x.split('->') for x in jkdescripts]
fjk = (ctypes.c_void_p*(n_dm))()
dmsptr = (ctypes.c_void_p*(n_dm))()
vjkptr = (ctypes.c_void_p*(n_dm))()
for i, (dmsym, vsym) in enumerate(descr_sym):
f1 = _fpointer('CVHFnr%s_%s_%s'%(aosym, dmsym, vsym))
assert(dms[i].shape == get_dims(dmsym, shls_slice, ao_loc))
vshape = (ncomp,) + get_dims(vsym[-2:], shls_slice, ao_loc)
vjk.append(numpy.empty(vshape))
dmsptr[i] = dms[i].ctypes.data_as(ctypes.c_void_p)
vjkptr[i] = vjk[i].ctypes.data_as(ctypes.c_void_p)
fjk[i] = f1
shls_slice = (ctypes.c_int*8)(*shls_slice)
fdrv(cintor, fdot, fjk, dmsptr, vjkptr,
ctypes.c_int(n_dm), ctypes.c_int(ncomp),
shls_slice, ao_loc.ctypes.data_as(ctypes.c_void_p), cintopt, cvhfopt,
c_atm.ctypes.data_as(ctypes.c_void_p), natm,
c_bas.ctypes.data_as(ctypes.c_void_p), nbas,
c_env.ctypes.data_as(ctypes.c_void_p))
if ncomp == 1:
vjk = [v.reshape(v.shape[1:]) for v in vjk]
if isinstance(jkdescript, str):
vjk = vjk[0]
return vjk
def eval_gto(mol, eval_name, coords,
comp=None, shls_slice=None, non0tab=None, ao_loc=None, out=None):
r'''Evaluate AO function value on the given grids,
Args:
eval_name : str
================== ====== =======================
Function comp Expression
================== ====== =======================
"GTOval_sph" 1 |AO>
"GTOval_ip_sph" 3 nabla |AO>
"GTOval_ig_sph" 3 (#C(0 1) g) |AO>
"GTOval_ipig_sph" 3 (#C(0 1) nabla g) |AO>
"GTOval_cart" 1 |AO>
"GTOval_ip_cart" 3 nabla |AO>
"GTOval_ig_cart" 3 (#C(0 1) g)|AO>
================== ====== =======================
atm : int32 ndarray
libcint integral function argument
bas : int32 ndarray
libcint integral function argument
env : float64 ndarray
libcint integral function argument
coords : 2D array, shape (N,3)
The coordinates of the grids.
Kwargs:
comp : int
Number of the components of the operator
shls_slice : 2-element list
(shl_start, shl_end).
If given, only part of AOs (shl_start <= shell_id < shl_end) are
evaluated. By default, all shells defined in mol will be evaluated.
non0tab : 2D bool array
mask array to indicate whether the AO values are zero. The mask
array can be obtained by calling :func:`dft.gen_grid.make_mask`
out : ndarray
If provided, results are written into this array.
Returns:
2D array of shape (N,nao) Or 3D array of shape (\*,N,nao) to store AO
values on grids.
Examples:
>>> mol = gto.M(atom='O 0 0 0; H 0 0 1; H 0 1 0', basis='ccpvdz')
>>> coords = numpy.random.random((100,3)) # 100 random points
>>> ao_value = mol.eval_gto("GTOval_sph", coords)
>>> print(ao_value.shape)
(100, 24)
>>> ao_value = mol.eval_gto("GTOval_ig_sph", coords)
>>> print(ao_value.shape)
(3, 100, 24)
'''
eval_name, comp = _get_intor_and_comp(mol, eval_name, comp)
atm = numpy.asarray(mol._atm, dtype=numpy.int32, order='C')
bas = numpy.asarray(mol._bas, dtype=numpy.int32, order='C')
env = numpy.asarray(mol._env, dtype=numpy.double, order='C')
coords = numpy.asarray(coords, dtype=numpy.double, order='F')
natm = atm.shape[0]
nbas = bas.shape[0]
ngrids = coords.shape[0]
if ao_loc is None:
ao_loc = make_loc(bas, eval_name)
if shls_slice is None:
shls_slice = (0, nbas)
sh0, sh1 = shls_slice
nao = ao_loc[sh1] - ao_loc[sh0]
if 'spinor' in eval_name:
ao = numpy.ndarray((2,comp,nao,ngrids), dtype=numpy.complex128, buffer=out)
else:
ao = numpy.ndarray((comp,nao,ngrids), buffer=out)
if non0tab is None:
non0tab = numpy.ones(((ngrids+BLKSIZE-1)//BLKSIZE,nbas),
dtype=numpy.int8)
drv = getattr(libcgto, eval_name)
drv(ctypes.c_int(ngrids),
(ctypes.c_int*2)(*shls_slice), ao_loc.ctypes.data_as(ctypes.c_void_p),
ao.ctypes.data_as(ctypes.c_void_p),
coords.ctypes.data_as(ctypes.c_void_p),
non0tab.ctypes.data_as(ctypes.c_void_p),
atm.ctypes.data_as(ctypes.c_void_p), ctypes.c_int(natm),
bas.ctypes.data_as(ctypes.c_void_p), ctypes.c_int(nbas),
env.ctypes.data_as(ctypes.c_void_p))
ao = numpy.swapaxes(ao, -1, -2)
if comp == 1:
if 'spinor' in eval_name:
ao = ao[:,0]
else:
ao = ao[0]
return ao
