def init_amps(mycc, eris=None):
eris = getattr(mycc, '_eris', None)
if eris is None:
mycc.ao2mo()
eris = mycc._eris
time0 = time.clock(), time.time()
mo_e = eris.mo_energy
nocc = mycc.nocc
nvir = mo_e.size - nocc
eia = mo_e[:nocc, None] - mo_e[None, nocc:]
t1T = eris.fock[nocc:, :nocc] / eia.T
loc0, loc1 = _task_location(nvir)
t2T = numpy.empty((loc1 - loc0, nvir, nocc, nocc))
max_memory = mycc.max_memory - lib.current_memory()[0]
blksize = int(
min(nvir, max(BLKMIN, max_memory * .3e6 / 8 / (nocc**2 * nvir + 1))))
emp2 = 0
for p0, p1 in lib.prange(0, loc1 - loc0, blksize):
eris_ovov = eris.ovov[:, p0:p1]
t2T[p0:p1] = (
eris_ovov.transpose(1, 3, 0, 2) /
lib.direct_sum('ia,jb->abij', eia[:, p0 + loc0:p1 + loc0], eia))
emp2 += 2 * numpy.einsum('abij,iajb', t2T[p0:p1], eris_ovov)
emp2 -= numpy.einsum('abji,iajb', t2T[p0:p1], eris_ovov)
mycc.emp2 = comm.allreduce(emp2)
logger.info(mycc, 'Init t2, MP2 energy = %.15g', mycc.emp2)
logger.timer(mycc, 'init mp2', *time0)
return mycc.emp2, t1T.T, t2T.transpose(2, 3, 0, 1)
def init_amps(mycc, eris=None):
eris = getattr(mycc, '_eris', None)
if eris is None:
mycc.ao2mo()
eris = mycc._eris
time0 = logger.process_clock(), logger.perf_counter()
mo_e = eris.mo_energy
nocc = mycc.nocc
nvir = mo_e.size - nocc
mo_e_o = mo_e[:nocc]
mo_e_v = mo_e[nocc:] + mycc.level_shift
eia = mo_e_o[:, None] - mo_e_v
t1T = eris.fock[nocc:, :nocc] / eia.T
loc0, loc1 = _task_location(nvir)
t2T = np.empty((loc1-loc0, nvir, nocc, nocc))
max_memory = mycc.max_memory - lib.current_memory()[0]
blksize = int(min(nvir, max(BLKMIN, max_memory*.3e6/8/(nocc**2*nvir+1))))
emp2 = 0
for p0, p1 in lib.prange(0, loc1-loc0, blksize):
eris_vvoo = eris.xvoo[p0:p1]
t2T[p0:p1] = (eris_vvoo / lib.direct_sum('ia, jb -> abij', eia[:, loc0+p0:loc0+p1], eia))
emp2 += np.einsum('abij, abij', t2T[p0:p1], eris_vvoo.conj(), optimize=True).real
eris_vvoo = None
mycc.emp2 = comm.allreduce(emp2) * 0.25
logger.info(mycc, 'Init t2, MP2 energy = %.15g', mycc.emp2)
logger.timer(mycc, 'init mp2', *time0)
mycc.t1 = t1T.T
mycc.t2 = t2T.transpose(2, 3, 0, 1)
return mycc.emp2, mycc.t1, mycc.t2
def init_amps(mycc, eris=None):
eris = getattr(mycc, '_eris', None)
if eris is None:
mycc.ao2mo()
eris = mycc._eris
time0 = time.clock(), time.time()
mo_e = eris.mo_energy
nocc = mycc.nocc
nvir = mo_e.size - nocc
eia = mo_e[:nocc,None] - mo_e[None,nocc:]
t1T = eris.fock[nocc:,:nocc] / eia.T
loc0, loc1 = _task_location(nvir)
t2T = numpy.empty((loc1-loc0,nvir,nocc,nocc))
max_memory = mycc.max_memory - lib.current_memory()[0]
blksize = int(min(nvir, max(BLKMIN, max_memory*.3e6/8/(nocc**2*nvir+1))))
emp2 = 0
for p0, p1 in lib.prange(0, loc1-loc0, blksize):
eris_ovov = eris.ovov[:,p0:p1]
t2T[p0:p1] = (eris_ovov.transpose(1,3,0,2) /
lib.direct_sum('ia,jb->abij', eia[:,p0+loc0:p1+loc0], eia))
emp2 += 2 * numpy.einsum('abij,iajb', t2T[p0:p1], eris_ovov)
emp2 -= numpy.einsum('abji,iajb', t2T[p0:p1], eris_ovov)
mycc.emp2 = comm.allreduce(emp2)
logger.info(mycc, 'Init t2, MP2 energy = %.15g', mycc.emp2)
logger.timer(mycc, 'init mp2', *time0)
return mycc.emp2, t1T.T, t2T.transpose(2,3,0,1)
def get_veff(mf, mol=None, dm=None, dm_last=0, vhf_last=0, hermi=1):
t0 = (time.clock(), time.time())
mf.unpack_(comm.bcast(mf.pack()))
mol = mf.mol
ni = mf._numint
if mf.nlc != '':
raise NotImplementedError
omega, alpha, hyb = ni.rsh_and_hybrid_coeff(mf.xc, spin=mol.spin)
if abs(omega) > 1e-10: # For range separated Coulomb operator
raise NotImplementedError
# Broadcast the large input arrays here.
if any(comm.allgather(isinstance(dm, str) and dm == 'SKIPPED_ARG')):
if rank == 0 and dm is None:
dm = mf.make_rdm1()
dm = mpi.bcast_tagged_array(dm)
if any(comm.allgather(isinstance(dm_last, str) and dm_last == 'SKIPPED_ARG')):
dm_last = mpi.bcast_tagged_array(dm_last)
if any(comm.allgather(isinstance(vhf_last, str) and vhf_last == 'SKIPPED_ARG')):
vhf_last = mpi.bcast_tagged_array(vhf_last)
ground_state = (isinstance(dm, numpy.ndarray) and dm.ndim == 2)
if mf.grids.coords is None:
_setup_grids_(mf, dm)
t0 = logger.timer(mf, 'setting up grids', *t0)
if hermi == 2: # because rho = 0
n, exc, vxc = 0, 0, 0
else:
n, exc, vxc = ni.nr_rks(mol, mf.grids, mf.xc, dm)
n = comm.allreduce(n)
exc = comm.allreduce(exc)
vxc = mpi.reduce(vxc)
logger.debug(mf, 'nelec by numeric integration = %s', n)
t0 = logger.timer(mf, 'vxc', *t0)
if abs(hyb) < 1e-10 and abs(alpha) < 1e-10:
vk = None
if getattr(vhf_last, 'vj', None) is not None:
ddm = numpy.asarray(dm) - dm_last
vj = mpi.reduce(mf.get_j(mol, ddm, hermi))
vj += vhf_last.vj
else:
vj = mf.get_j(mol, dm, hermi)
vj = mpi.reduce(vj)
vxc += vj
else:
if getattr(vhf_last, 'vk', None) is not None:
ddm = numpy.asarray(dm) - dm_last
vj, vk = mf.get_jk(mol, ddm, hermi)
ddm = None
vj = mpi.reduce(vj)
vk = mpi.reduce(vk) * hyb
vj += vhf_last.vj
vk += vhf_last.vk
else:
vj, vk = mf.get_jk(mol, dm, hermi)
vj = mpi.reduce(vj)
vk = mpi.reduce(vk) * hyb
vxc += vj - vk * .5
if ground_state:
exc -= numpy.einsum('ij,ji', dm, vk) * .5 * .5
if ground_state:
ecoul = numpy.einsum('ij,ji', dm, vj) * .5
else:
ecoul = None
vxc = lib.tag_array(vxc, ecoul=ecoul, exc=exc, vj=vj, vk=vk)
return vxc
def get_veff(mf, mol=None, dm=None, dm_last=0, vhf_last=0, hermi=1):
t0 = (logger.process_clock(), logger.perf_counter())
mf.unpack_(comm.bcast(mf.pack()))
mol = mf.mol
ni = mf._numint
if mf.nlc != '':
raise NotImplementedError
omega, alpha, hyb = ni.rsh_and_hybrid_coeff(mf.xc, spin=mol.spin)
# Broadcast the large input arrays here.
if any(comm.allgather(dm is mpi.Message.SkippedArg)):
if rank == 0 and dm is None:
dm = mf.make_rdm1()
dm = mpi.bcast_tagged_array(dm)
if any(comm.allgather(dm_last is mpi.Message.SkippedArg)):
dm_last = mpi.bcast_tagged_array(dm_last)
if any(comm.allgather(vhf_last is mpi.Message.SkippedArg)):
vhf_last = mpi.bcast_tagged_array(vhf_last)
ground_state = (dm.ndim == 3 and dm.shape[0] == 2)
if mf.grids.coords is None:
mpi_rks._setup_grids_(mf, dm[0]+dm[1])
t0 = logger.timer(mf, 'setting up grids', *t0)
if hermi == 2: # because rho = 0
n, exc, vxc = 0, 0, 0
else:
n, exc, vxc = ni.nr_uks(mol, mf.grids, mf.xc, dm)
n = comm.allreduce(n)
exc = comm.allreduce(exc)
vxc = mpi.reduce(vxc)
logger.debug(mf, 'nelec by numeric integration = %s', n)
t0 = logger.timer(mf, 'vxc', *t0)
if abs(hyb) < 1e-10 and abs(alpha) < 1e-10:
vk = None
if getattr(vhf_last, 'vj', None) is not None:
ddm = numpy.asarray(dm) - dm_last
ddm = ddm[0] + ddm[1]
vj = mf.get_j(mol, ddm, hermi)
vj += vhf_last.vj
else:
vj = mf.get_j(mol, dm[0]+dm[1], hermi)
vxc += vj
else:
if getattr(vhf_last, 'vk', None) is not None:
ddm = numpy.asarray(dm) - dm_last
vj, vk = mf.get_jk(mol, ddm, hermi)
vj = vj[0] + vj[1]
vk *= hyb
if abs(omega) > 1e-10:
vklr = mf.get_k(mol, ddm, hermi, omega=omega)
vk += vklr * (alpha - hyb)
ddm = None
vj += vhf_last.vj
vk += vhf_last.vk
else:
vj, vk = mf.get_jk(mol, dm, hermi)
vj = vj[0] + vj[1]
vk *= hyb
if abs(omega) > 1e-10:
vklr = mf.get_k(mol, dm, hermi, omega=omega)
vk += vklr * (alpha - hyb)
vxc += vj
vxc -= vk
if ground_state:
exc -=(numpy.einsum('ij,ji', dm[0], vk[0]) +
numpy.einsum('ij,ji', dm[1], vk[1])) * .5
if ground_state:
ecoul = numpy.einsum('ij,ji', dm[0]+dm[1], vj) * .5
else:
ecoul = None
vxc = lib.tag_array(vxc, ecoul=ecoul, exc=exc, vj=vj, vk=vk)
return vxc
def _eval_jk(mf, dm, hermi, gen_jobs):
cpu0 = (logger.process_clock(), logger.perf_counter())
mol = mf.mol
ao_loc = mol.ao_loc_nr()
nao = ao_loc[-1]
bas_groups = _partition_bas(mol)
jobs = gen_jobs(len(bas_groups), hermi)
njobs = len(jobs)
logger.debug1(mf, 'njobs %d', njobs)
# Each job has multiple recipes.
n_recipes = len(jobs[0][1:])
dm = numpy.asarray(dm).reshape(-1, nao, nao)
n_dm = dm.shape[0]
vk = numpy.zeros((n_recipes, n_dm, nao, nao))
if mf.opt is None:
vhfopt = mf.init_direct_scf(mol)
else:
vhfopt = mf.opt
# Assign the entire dm_cond to vhfopt.
# The prescreen function CVHFnrs8_prescreen will index q_cond and dm_cond
# over the entire basis. "set_dm" in function jk.get_jk/direct_bindm only
# creates a subblock of dm_cond which is not compatible with
# CVHFnrs8_prescreen.
vhfopt.set_dm(dm, mol._atm, mol._bas, mol._env)
# Then skip the "set_dm" initialization in function jk.get_jk/direct_bindm.
vhfopt._dmcondname = None
logger.timer_debug1(mf, 'get_jk initialization', *cpu0)
for job_id in mpi.work_stealing_partition(range(njobs)):
group_ids = jobs[job_id][0]
recipes = jobs[job_id][1:]
shls_slice = lib.flatten([bas_groups[i] for i in group_ids])
loc = ao_loc[shls_slice].reshape(4, 2)
dm_blks = []
for i_dm in range(n_dm):
for ir, recipe in enumerate(recipes):
for i, rec in enumerate(recipe):
p0, p1 = loc[rec[0]]
q0, q1 = loc[rec[1]]
dm_blks.append(dm[i_dm, p0:p1, q0:q1])
scripts = [
'ijkl,%s%s->%s%s' % tuple(['ijkl'[x] for x in rec])
for recipe in recipes for rec in recipe
] * n_dm
kparts = jk.get_jk(mol,
dm_blks,
scripts,
shls_slice=shls_slice,
vhfopt=vhfopt)
for i_dm in range(n_dm):
for ir, recipe in enumerate(recipes):
for i, rec in enumerate(recipe):
p0, p1 = loc[rec[2]]
q0, q1 = loc[rec[3]]
vk[ir, i_dm, p0:p1, q0:q1] += kparts[i]
# Pop the results of one recipe
kparts = kparts[i + 1:]
vk = mpi.reduce(vk)
if rank == 0:
if hermi:
for i in range(n_recipes):
for j in range(n_dm):
lib.hermi_triu(vk[i, j], hermi, inplace=True)
else:
# Zero out vk on workers. If reduce(get_jk()) is called twice,
# non-zero vk on workers can cause error.
vk[:] = 0
logger.timer(mf, 'get_jk', *cpu0)
return vk