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 _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