def ao2mo(self, mo_coeff=None, store_eris=False):
eris = mp2._ChemistsERIs()
# Initialize only the mo_coeff
eris._common_init_(self, mo_coeff)
if store_eris:
nocc = self.nocc
nvir = self.nmo - nocc
naux = self.with_df.get_naoaux()
Lov = numpy.empty((naux, nocc*nvir))
p1 = 0
for istep, qov in enumerate(self.loop_ao2mo(eris.mo_coeff, nocc)):
logger.debug(self, 'Load cderi step %d', istep)
p0, p1 = p1, p1 + qov.shape[0]
Lov[p0:p1] = qov
Lov = Lov.reshape(naux, nocc, nvir)
eris.ovov = numpy.einsum("Qia,Qjb->iajb", Lov, Lov, optimize=True)
return eris
def ao2mo(self, mo_coeff=None):
eris = mp2._ChemistsERIs()
# Initialize only the mo_coeff and
eris._common_init_(self, mo_coeff)
return eris
def _make_eris(mp, mo_coeff=None, verbose=None):
log = logger.new_logger(mp, verbose)
time0 = (time.clock(), time.time())
log.debug('transform (ia|jb) outcore')
mol = mp.mol
nocc = mp.nocc
nmo = mp.nmo
nvir = nmo - nocc
eris = mp2._ChemistsERIs(mp, mo_coeff)
nao = eris.mo_coeff.shape[0]
assert (nvir <= nao)
orbo = eris.mo_coeff[:, :nocc]
orbv = numpy.asarray(eris.mo_coeff[:, nocc:], order='F')
eris.feri = lib.H5TmpFile()
int2e = mol._add_suffix('int2e')
ao2mopt = _ao2mo.AO2MOpt(mol, int2e, 'CVHFnr_schwarz_cond',
'CVHFsetnr_direct_scf')
fint = gto.moleintor.getints4c
ntasks = mpi.pool.size
olocs = [_task_location(nocc, task_id) for task_id in range(ntasks)]
oloc0, oloc1 = olocs[rank]
nocc_seg = oloc1 - oloc0
log.debug2('olocs %s', olocs)
ao_loc = mol.ao_loc_nr()
task_sh_locs = lib.misc._balanced_partition(ao_loc, ntasks)
log.debug2('task_sh_locs %s', task_sh_locs)
ao_sh0 = task_sh_locs[rank]
ao_sh1 = task_sh_locs[rank + 1]
ao_loc0 = ao_loc[ao_sh0]
ao_loc1 = ao_loc[ao_sh1]
nao_seg = ao_loc1 - ao_loc0
orbo_seg = orbo[ao_loc0:ao_loc1]
mem_now = lib.current_memory()[0]
max_memory = max(0, mp.max_memory - mem_now)
dmax = numpy.sqrt(max_memory * .9e6 / 8 / ((nao + nocc) *
(nao_seg + nocc)))
dmax = min(nao // 4 + 2, max(BLKMIN, min(comm.allgather(dmax))))
sh_ranges = ao2mo.outcore.balance_partition(ao_loc, dmax)
sh_ranges = comm.bcast(sh_ranges)
dmax = max(x[2] for x in sh_ranges)
eribuf = numpy.empty((nao, dmax, dmax, nao_seg))
ftmp = lib.H5TmpFile()
log.debug('max_memory %s MB (dmax = %s) required disk space %g MB',
max_memory, dmax,
nocc * nocc_seg * (nao * (nao + dmax) / 2 + nvir**2) * 8 / 1e6)
def save(count, tmp_xo):
di, dj = tmp_xo.shape[2:4]
tmp_xo = [tmp_xo[p0:p1] for p0, p1 in olocs]
tmp_xo = mpi.alltoall(tmp_xo, split_recvbuf=True)
tmp_xo = sum(tmp_xo).reshape(nocc_seg, nocc, di, dj)
ftmp[str(count) + 'b'] = tmp_xo
tmp_ox = mpi.alltoall([tmp_xo[:, p0:p1] for p0, p1 in olocs],
split_recvbuf=True)
tmp_ox = [
tmp_ox[i].reshape(p1 - p0, nocc_seg, di, dj)
for i, (p0, p1) in enumerate(olocs)
]
ftmp[str(count) + 'a'] = numpy.vstack(tmp_ox)
jk_blk_slices = []
count = 0
time1 = time0
with lib.call_in_background(save) as bg_save:
for ip, (ish0, ish1, ni) in enumerate(sh_ranges):
for jsh0, jsh1, nj in sh_ranges[:ip + 1]:
i0, i1 = ao_loc[ish0], ao_loc[ish1]
j0, j1 = ao_loc[jsh0], ao_loc[jsh1]
jk_blk_slices.append((i0, i1, j0, j1))
shls_slice = (0, mol.nbas, ish0, ish1, jsh0, jsh1, ao_sh0,
ao_sh1)
eri = fint(int2e,
mol._atm,
mol._bas,
mol._env,
shls_slice=shls_slice,
aosym='s1',
ao_loc=ao_loc,
cintopt=ao2mopt._cintopt,
out=eribuf)
tmp_xo = lib.einsum('pi,pqrs->iqrs', orbo, eri)
tmp_xo = lib.einsum('iqrs,sl->ilqr', tmp_xo, orbo_seg)
bg_save(count, tmp_xo)
tmp_xo = None
count += 1
time1 = log.timer_debug1(
'partial ao2mo [%d:%d,%d:%d]' % (ish0, ish1, jsh0, jsh1),
*time1)
eri = eribuf = None
time1 = time0 = log.timer('mp2 ao2mo_ovov pass1', *time0)
eris.ovov = eris.feri.create_dataset('ovov', (nocc, nvir, nocc_seg, nvir),
'f8')
occblk = int(
min(nocc,
max(BLKMIN, max_memory * .9e6 / 8 / (nao**2 * nocc_seg + 1) / 5)))
def load(i0, eri):
if i0 < nocc:
i1 = min(i0 + occblk, nocc)
for k, (p0, p1, q0, q1) in enumerate(jk_blk_slices):
eri[:i1 - i0, :, p0:p1, q0:q1] = ftmp[str(k) + 'a'][i0:i1]
if p0 != q0:
dat = numpy.asarray(ftmp[str(k) + 'b'][:, i0:i1])
eri[:i1 - i0, :, q0:q1, p0:p1] = dat.transpose(1, 0, 3, 2)
def save(i0, i1, dat):
eris.ovov[i0:i1] = dat
buf_prefecth = numpy.empty((occblk, nocc_seg, nao, nao))
buf = numpy.empty_like(buf_prefecth)
bufw = numpy.empty((occblk * nocc_seg, nvir**2))
bufw1 = numpy.empty_like(bufw)
with lib.call_in_background(load) as prefetch:
with lib.call_in_background(save) as bsave:
load(0, buf_prefecth)
for i0, i1 in lib.prange(0, nocc, occblk):
buf, buf_prefecth = buf_prefecth, buf
prefetch(i1, buf_prefecth)
eri = buf[:i1 - i0].reshape((i1 - i0) * nocc_seg, nao, nao)
dat = _ao2mo.nr_e2(eri,
orbv, (0, nvir, 0, nvir),
's1',
's1',
out=bufw)
bsave(
i0, i1,
dat.reshape(i1 - i0, nocc_seg, nvir,
nvir).transpose(0, 2, 1, 3))
bufw, bufw1 = bufw1, bufw
time1 = log.timer_debug1('pass2 ao2mo [%d:%d]' % (i0, i1),
*time1)
time0 = log.timer('mp2 ao2mo_ovov pass2', *time0)
mp._eris = eris
return eris
def _make_eris(mp, mo_coeff=None, verbose=None):
log = logger.new_logger(mp, verbose)
time0 = (time.clock(), time.time())
log.debug('transform (ia|jb) outcore')
mol = mp.mol
nocc = mp.nocc
nmo = mp.nmo
nvir = nmo - nocc
eris = mp2._ChemistsERIs(mp, mo_coeff)
nao = eris.mo_coeff.shape[0]
assert(nvir <= nao)
orbo = eris.mo_coeff[:,:nocc]
orbv = numpy.asarray(eris.mo_coeff[:,nocc:], order='F')
eris.feri = lib.H5TmpFile()
int2e = mol._add_suffix('int2e')
ao2mopt = _ao2mo.AO2MOpt(mol, int2e, 'CVHFnr_schwarz_cond',
'CVHFsetnr_direct_scf')
fint = gto.moleintor.getints4c
ntasks = mpi.pool.size
olocs = [_task_location(nocc, task_id) for task_id in range(ntasks)]
oloc0, oloc1 = olocs[rank]
nocc_seg = oloc1 - oloc0
log.debug2('olocs %s', olocs)
ao_loc = mol.ao_loc_nr()
task_sh_locs = lib.misc._balanced_partition(ao_loc, ntasks)
log.debug2('task_sh_locs %s', task_sh_locs)
ao_sh0 = task_sh_locs[rank]
ao_sh1 = task_sh_locs[rank+1]
ao_loc0 = ao_loc[ao_sh0]
ao_loc1 = ao_loc[ao_sh1]
nao_seg = ao_loc1 - ao_loc0
orbo_seg = orbo[ao_loc0:ao_loc1]
mem_now = lib.current_memory()[0]
max_memory = max(0, mp.max_memory - mem_now)
dmax = numpy.sqrt(max_memory*.9e6/8/((nao+nocc)*(nao_seg+nocc)))
dmax = min(nao//4+2, max(BLKMIN, min(comm.allgather(dmax))))
sh_ranges = ao2mo.outcore.balance_partition(ao_loc, dmax)
sh_ranges = comm.bcast(sh_ranges)
dmax = max(x[2] for x in sh_ranges)
eribuf = numpy.empty((nao,dmax,dmax,nao_seg))
ftmp = lib.H5TmpFile()
log.debug('max_memory %s MB (dmax = %s) required disk space %g MB',
max_memory, dmax, nocc*nocc_seg*(nao*(nao+dmax)/2+nvir**2)*8/1e6)
def save(count, tmp_xo):
di, dj = tmp_xo.shape[2:4]
tmp_xo = [tmp_xo[p0:p1] for p0, p1 in olocs]
tmp_xo = mpi.alltoall(tmp_xo, split_recvbuf=True)
tmp_xo = sum(tmp_xo).reshape(nocc_seg,nocc,di,dj)
ftmp[str(count)+'b'] = tmp_xo
tmp_ox = mpi.alltoall([tmp_xo[:,p0:p1] for p0, p1 in olocs],
split_recvbuf=True)
tmp_ox = [tmp_ox[i].reshape(p1-p0,nocc_seg,di,dj)
for i, (p0,p1) in enumerate(olocs)]
ftmp[str(count)+'a'] = numpy.vstack(tmp_ox)
jk_blk_slices = []
count = 0
time1 = time0
with lib.call_in_background(save) as bg_save:
for ip, (ish0, ish1, ni) in enumerate(sh_ranges):
for jsh0, jsh1, nj in sh_ranges[:ip+1]:
i0, i1 = ao_loc[ish0], ao_loc[ish1]
j0, j1 = ao_loc[jsh0], ao_loc[jsh1]
jk_blk_slices.append((i0,i1,j0,j1))
shls_slice = (0,mol.nbas,ish0,ish1, jsh0,jsh1,ao_sh0,ao_sh1)
eri = fint(int2e, mol._atm, mol._bas, mol._env,
shls_slice=shls_slice, aosym='s1', ao_loc=ao_loc,
cintopt=ao2mopt._cintopt, out=eribuf)
tmp_xo = lib.einsum('pi,pqrs->iqrs', orbo, eri)
tmp_xo = lib.einsum('iqrs,sl->ilqr', tmp_xo, orbo_seg)
bg_save(count, tmp_xo)
tmp_xo = None
count += 1
time1 = log.timer_debug1('partial ao2mo [%d:%d,%d:%d]' %
(ish0,ish1,jsh0,jsh1), *time1)
eri = eribuf = None
time1 = time0 = log.timer('mp2 ao2mo_ovov pass1', *time0)
eris.ovov = eris.feri.create_dataset('ovov', (nocc,nvir,nocc_seg,nvir), 'f8')
occblk = int(min(nocc, max(BLKMIN, max_memory*.9e6/8/(nao**2*nocc_seg+1)/5)))
def load(i0, eri):
if i0 < nocc:
i1 = min(i0+occblk, nocc)
for k, (p0,p1,q0,q1) in enumerate(jk_blk_slices):
eri[:i1-i0,:,p0:p1,q0:q1] = ftmp[str(k)+'a'][i0:i1]
if p0 != q0:
dat = numpy.asarray(ftmp[str(k)+'b'][:,i0:i1])
eri[:i1-i0,:,q0:q1,p0:p1] = dat.transpose(1,0,3,2)
def save(i0, i1, dat):
eris.ovov[i0:i1] = dat
buf_prefecth = numpy.empty((occblk,nocc_seg,nao,nao))
buf = numpy.empty_like(buf_prefecth)
bufw = numpy.empty((occblk*nocc_seg,nvir**2))
bufw1 = numpy.empty_like(bufw)
with lib.call_in_background(load) as prefetch:
with lib.call_in_background(save) as bsave:
load(0, buf_prefecth)
for i0, i1 in lib.prange(0, nocc, occblk):
buf, buf_prefecth = buf_prefecth, buf
prefetch(i1, buf_prefecth)
eri = buf[:i1-i0].reshape((i1-i0)*nocc_seg,nao,nao)
dat = _ao2mo.nr_e2(eri, orbv, (0,nvir,0,nvir), 's1', 's1', out=bufw)
bsave(i0, i1, dat.reshape(i1-i0,nocc_seg,nvir,nvir).transpose(0,2,1,3))
bufw, bufw1 = bufw1, bufw
time1 = log.timer_debug1('pass2 ao2mo [%d:%d]' % (i0,i1), *time1)
time0 = log.timer('mp2 ao2mo_ovov pass2', *time0)
mp._eris = eris
return eris