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.CVHFnr_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))
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":
fsetqcond = getattr(libcvhf, qcondname)
fsetqcond(
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),
)
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 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')
out = numpy.ndarray((comp,nkl,ij_count), dtype=numpy.complex, buffer=out)
if out.size == 0:
return out
if ao2mopt is not None:
cao2mopt = ao2mopt._this
cintopt = ao2mopt._cintopt
intor = ao2mopt._intor
else:
cao2mopt = lib.c_null_ptr()
cintopt = make_cintopt(c_atm, c_bas, c_env, intor)
cintor = _fpointer(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 nr_e1fill(intor, sh_range, atm, bas, env,
aosym='s1', comp=1, ao2mopt=None, out=None):
assert(aosym in ('s4', 's2ij', 's2kl', 's1'))
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, order='C')
natm = ctypes.c_int(c_atm.shape[0])
nbas = ctypes.c_int(c_bas.shape[0])
ao_loc = 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
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
intor = ao2mopt._intor
else:
cao2mopt = lib.c_null_ptr()
cintopt = make_cintopt(c_atm, c_bas, c_env, intor)
cintor = _fpointer(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 _int3c_wrapper(mol, auxmol, intor, aosym):
nbas = mol.nbas
pmol = mol + auxmol
intor = mol._add_suffix(intor)
opt = make_cintopt(mol._atm, mol._bas, mol._env, intor)
def get_int3c(shls_slice=None):
if shls_slice is None:
shls_slice = (0, nbas, 0, nbas, nbas, pmol.nbas)
else:
shls_slice = shls_slice[:4] + (nbas + shls_slice[4],
nbas + shls_slice[5])
return getints(intor,
pmol._atm,
pmol._bas,
pmol._env,
shls_slice,
aosym=aosym,
cintopt=opt)
return get_int3c
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 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 direct_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")
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 = 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 = numpy.asarray(make_ao_loc(bas), dtype=numpy.int32)
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:
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 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 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 rdirect_bindm(intor, aosym, jkdescript, dms, ncomp, atm, bas, env, vhfopt=None):
assert aosym in ("s8", "s4", "s2ij", "s2kl", "s1", "a4ij", "a4kl", "a2ij", "a2kl")
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 = vhfopt._intor
fdrv = getattr(libcvhf, "CVHFr_direct_drv")
dotsym = _INTSYMAP[aosym]
fdot = _fpointer("CVHFdot_r" + dotsym)
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),
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):
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 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_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 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):
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:
cintor = _fpointer("cint2e_sph")
cintopt = make_cintopt(c_atm, c_bas, c_env, "cint2e_sph")
cvhfopt = pyscf.lib.c_null_ptr()
else:
vhfopt.set_dm(dms, atm, bas, env)
cvhfopt = vhfopt._this
cintopt = vhfopt._cintopt
cintor = vhfopt._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 = numpy.asarray(make_ao_loc(bas), dtype=numpy.int32)
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