def extrapolate(self, nd=None):
if nd is None:
nd = self.get_num_vec()
if nd == 0:
raise RuntimeError('No vector found in DIIS object.')
h = self._H[:nd+1,:nd+1].copy()
h[1:,1:] = mpi.comm.allreduce(self._H[1:nd+1,1:nd+1])
g = numpy.zeros(nd+1, h.dtype)
g[0] = 1
w, v = scipy.linalg.eigh(h)
if numpy.any(abs(w)<1e-14):
logger.debug(self, 'Singularity found in DIIS error vector space.')
idx = abs(w)>1e-14
c = numpy.dot(v[:,idx]*(1./w[idx]), numpy.dot(v[:,idx].T.conj(), g))
else:
try:
c = numpy.linalg.solve(h, g)
except numpy.linalg.linalg.LinAlgError as e:
logger.warn(self, ' diis singular, eigh(h) %s', w)
raise e
logger.debug1(self, 'diis-c %s', c)
xnew = None
for i, ci in enumerate(c[1:]):
xi = self.get_vec(i)
if xnew is None:
xnew = numpy.zeros(xi.size, c.dtype)
for p0, p1 in lib.prange(0, xi.size, lib.diis.BLOCK_SIZE):
xnew[p0:p1] += xi[p0:p1] * ci
return xnew
def extrapolate(self, nd=None):
if nd is None:
nd = self.get_num_vec()
if nd == 0:
raise RuntimeError('No vector found in DIIS object.')
h = self._H[:nd + 1, :nd + 1].copy()
h[1:, 1:] = mpi.comm.allreduce(self._H[1:nd + 1, 1:nd + 1])
g = numpy.zeros(nd + 1, h.dtype)
g[0] = 1
w, v = scipy.linalg.eigh(h)
if numpy.any(abs(w) < 1e-14):
logger.debug(self, 'Singularity found in DIIS error vector space.')
idx = abs(w) > 1e-14
c = numpy.dot(v[:, idx] * (1. / w[idx]),
numpy.dot(v[:, idx].T.conj(), g))
else:
try:
c = numpy.linalg.solve(h, g)
except numpy.linalg.linalg.LinAlgError as e:
logger.warn(self, ' diis singular, eigh(h) %s', w)
raise e
logger.debug1(self, 'diis-c %s', c)
xnew = None
for i, ci in enumerate(c[1:]):
xi = self.get_vec(i)
if xnew is None:
xnew = numpy.zeros(xi.size, c.dtype)
for p0, p1 in lib.prange(0, xi.size, lib.diis.BLOCK_SIZE):
xnew[p0:p1] += xi[p0:p1] * ci
return xnew
def run_diis(self, t1, t2, istep, normt, de, adiis):
if (adiis and istep >= self.diis_start_cycle
and abs(de) < self.diis_start_energy_diff):
vec = self.amplitudes_to_vector(t1, t2)
t1, t2 = self.vector_to_amplitudes(adiis.update(vec))
logger.debug1(self, 'DIIS for step %d', istep)
return t1, t2
def run_diis(self, t1, t2, istep, normt, de, adiis):
if (adiis and
istep >= self.diis_start_cycle and
abs(de) < self.diis_start_energy_diff):
vec = self.amplitudes_to_vector(t1, t2)
t1, t2 = self.vector_to_amplitudes(adiis.update(vec))
logger.debug1(self, 'DIIS for step %d', istep)
return t1, t2
def _partition_bas(mol):
ao_loc = mol.ao_loc_nr()
nao = ao_loc[-1]
ngroups = max((mpi.pool.size * 50 * 8)**.25, 9)
blksize = max(BLKSIZE_MIN, min(nao / ngroups, BLKSIZE_MAX))
groups = ao2mo.outcore.balance_partition(ao_loc, blksize)
bas_groups = [x[:2] for x in groups]
logger.debug1(mol, 'mpi.size %d, blksize = %d, ngroups = %d',
mpi.pool.size, blksize, len(bas_groups))
logger.debug2(mol, 'bas_groups = %s', bas_groups)
return bas_groups
def __init__(self, mycc):
self._cc = mycc
self.daemon = None
nocc, nvir = mycc.t1.shape
nmo = nocc + nvir
nvir_seg = (nvir + mpi.pool.size - 1) // mpi.pool.size
max_memory = mycc.max_memory - lib.current_memory()[0]
max_memory = max(0, max_memory - nocc**3*13*lib.num_threads()*8/1e6)
blksize = max(BLKMIN, (max_memory*.5e6/8/(6*nmo*nocc))**.5 - nocc/4)
blksize = int(min(comm.allgather(min(nvir/6+2, nvir_seg/2+1, blksize))))
logger.debug1(mycc, 'GlobalDataHandler blksize %s', blksize)
self.vranges = []
self.data_partition = []
self.segment_location = {}
for task in range(mpi.pool.size):
p0 = nvir_seg * task
p1 = min(nvir, p0 + nvir_seg)
self.vranges.append((p0, p1))
for j0, j1 in lib.prange(p0, p1, blksize):
self.data_partition.append((j0, j1))
self.segment_location[j0] = task
logger.debug1(mycc, 'data_partition %s', self.data_partition)
logger.debug1(mycc, 'segment_location %s', self.segment_location)
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
