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