def _ztrans(Lpq, zij, moij, ijslice, Lrs, zkl, mokl, klslice, sym): tao = [] ao_loc = None zij = _ao2mo.r_e2(Lpq, moij, ijslice, tao, ao_loc, out=zij) if sym: zkl = zij else: zkl = _ao2mo.r_e2(Lrs, mokl, klslice, tao, ao_loc, out=zkl) return zij, zkl
def general(mydf, mo_coeffs, kpts=None, compact=True): if isinstance(mo_coeffs, numpy.ndarray) and mo_coeffs.ndim == 2: mo_coeffs = (mo_coeffs,) * 4 eri = mydf.get_eri(kpts) #################### # gamma point, the integral is real and with s4 symmetry if eri.dtype == numpy.float64: return ao2mo.general(eri, mo_coeffs, compact=compact) else: mokl, klslice = ao2mo.incore._conc_mos(mo_coeffs[2], mo_coeffs[3], False)[2:] if mokl.dtype == numpy.float64: mokl = mokl + 0j nao = mo_coeffs[0].shape[0] nmoi = mo_coeffs[0].shape[1] nmoj = mo_coeffs[1].shape[1] nmok = mo_coeffs[2].shape[1] nmol = mo_coeffs[3].shape[1] moi = numpy.asarray(mo_coeffs[0], order='F') moj = numpy.asarray(mo_coeffs[1], order='F') tao = [0] ao_loc = None pqkl = _ao2mo.r_e2(eri.reshape(-1,nao**2), mokl, klslice, tao, ao_loc, aosym='s1') pqkl = pqkl.reshape(nao,nao,nmok*nmol) pjkl = numpy.empty((nao,nmoj,nmok*nmol), dtype=numpy.complex128) for i in range(nao): lib.dot(moj.T, pqkl[i], 1, pjkl[i], 0) pqkl = None eri_mo = lib.dot(moi.T.conj(), pjkl.reshape(nao,-1)) return eri_mo.reshape(nmoi*nmoj,-1)
def _init_df_eris(cc, eris): from pyscf.pbc.df import df from pyscf.ao2mo import _ao2mo if cc._scf.with_df._cderi is None: cc._scf.with_df.build() cell = cc._scf.cell if cell.dimension == 2: # 2D ERIs are not positive definite. The 3-index tensors are stored in # two part. One corresponds to the positive part and one corresponds # to the negative part. The negative part is not considered in the # DF-driven CCSD implementation. raise NotImplementedError nocc = cc.nocc nmo = cc.nmo nvir = nmo - nocc nao = cell.nao_nr() kpts = cc.kpts nkpts = len(kpts) naux = cc._scf.with_df.get_naoaux() if gamma_point(kpts): dtype = np.double else: dtype = np.complex128 dtype = np.result_type(dtype, *eris.mo_coeff) eris.Lpv = Lpv = np.empty((nkpts, nkpts), dtype=object) with h5py.File(cc._scf.with_df._cderi, 'r') as f: kptij_lst = f['j3c-kptij'].value tao = [] ao_loc = None for ki, kpti in enumerate(kpts): for kj, kptj in enumerate(kpts): kpti_kptj = np.array((kpti, kptj)) Lpq = np.asarray(df._getitem(f, 'j3c', kpti_kptj, kptij_lst)) mo = np.hstack((eris.mo_coeff[ki], eris.mo_coeff[kj][:, nocc:])) mo = np.asarray(mo, dtype=dtype, order='F') if dtype == np.double: out = _ao2mo.nr_e2(Lpq, mo, (0, nmo, nmo, nmo + nvir), aosym='s2') else: #Note: Lpq.shape[0] != naux if linear dependency is found in auxbasis if Lpq[0].size != nao**2: # aosym = 's2' Lpq = lib.unpack_tril(Lpq).astype(np.complex128) out = _ao2mo.r_e2(Lpq, mo, (0, nmo, nmo, nmo + nvir), tao, ao_loc) Lpv[ki, kj] = out.reshape(-1, nmo, nvir) return eris
def general(mydf, mo_coeffs, kpts=None, compact=True): if isinstance(mo_coeffs, numpy.ndarray) and mo_coeffs.ndim == 2: mo_coeffs = (mo_coeffs, ) * 4 eri = mydf.get_eri(kpts) #################### # gamma point, the integral is real and with s4 symmetry if eri.dtype == numpy.float64: return ao2mo.general(eri, mo_coeffs, compact=compact) else: mokl, klslice = ao2mo.incore._conc_mos(mo_coeffs[2], mo_coeffs[3], False)[2:] if mokl.dtype == numpy.float64: mokl = mokl + 0j nao = mo_coeffs[0].shape[0] nmoi = mo_coeffs[0].shape[1] nmoj = mo_coeffs[1].shape[1] nmok = mo_coeffs[2].shape[1] nmol = mo_coeffs[3].shape[1] moi = numpy.asarray(mo_coeffs[0], order='F') moj = numpy.asarray(mo_coeffs[1], order='F') tao = [0] ao_loc = None pqkl = _ao2mo.r_e2(eri.reshape(-1, nao**2), mokl, klslice, tao, ao_loc, aosym='s1') pqkl = pqkl.reshape(nao, nao, nmok * nmol) pjkl = numpy.empty((nao, nmoj, nmok * nmol), dtype=numpy.complex128) for i in range(nao): lib.dot(moj.T, pqkl[i], 1, pjkl[i], 0) pqkl = None eri_mo = lib.dot(moi.T.conj(), pjkl.reshape(nao, -1)) return eri_mo.reshape(nmoi * nmoj, -1)
def get_sigmaR_diag(gw, omega, kn, orbp, ef, freqs, qij, q_abs): ''' Compute self-energy for poles inside coutour (more and more expensive away from Fermi surface) ''' mo_energy = np.array(gw._scf.mo_energy) mo_coeff = np.array(gw._scf.mo_coeff) nocc = gw.nocc nmo = gw.nmo nkpts = gw.nkpts kpts = gw.kpts nw = len(freqs) mydf = gw.with_df # possible kpts shift center kscaled = gw.mol.get_scaled_kpts(kpts) kscaled -= kscaled[0] idx = [] for k in range(nkpts): if omega > ef: fm = 1.0 idx.append( np.where((mo_energy[k] < omega) & (mo_energy[k] > ef))[0]) else: fm = -1.0 idx.append( np.where((mo_energy[k] > omega) & (mo_energy[k] < ef))[0]) sigmaR = 0j for kL in range(nkpts): # Lij: (ki, L, i, j) for looping every kL Lij = [] # kidx: save kj that conserves with kL and ki (-ki+kj+kL=G) # kidx_r: save ki that conserves with kL and kj (-ki+kj+kL=G) kidx = np.zeros((nkpts), dtype=np.int64) kidx_r = np.zeros((nkpts), dtype=np.int64) for i, kpti in enumerate(kpts): for j, kptj in enumerate(kpts): # Find (ki,kj) that satisfies momentum conservation with kL kconserv = -kscaled[i] + kscaled[j] + kscaled[kL] is_kconserv = np.linalg.norm(np.round(kconserv) - kconserv) < 1e-12 if is_kconserv: kidx[i] = j kidx_r[j] = i km = kidx_r[kn] if len(idx[km]) > 0: for i, kpti in enumerate(kpts): for j, kptj in enumerate(kpts): # Find (ki,kj) that satisfies momentum conservation with kL kconserv = -kscaled[i] + kscaled[j] + kscaled[kL] is_kconserv = np.linalg.norm( np.round(kconserv) - kconserv) < 1e-12 if is_kconserv: kidx[i] = j kidx_r[j] = i #logger.debug(gw, "Read Lpq (kL: %s / %s, ki: %s, kj: %s)"%(kL+1, nkpts, i, j)) Lij_out = None # Read (L|pq) and ao2mo transform to (L|ij) Lpq = [] for LpqR, LpqI, sign in mydf.sr_loop( [kpti, kptj], max_memory=0.1 * gw._scf.max_memory, compact=False): Lpq.append(LpqR + LpqI * 1.0j) # support uneqaul naux on different k points Lpq = np.vstack(Lpq).reshape(-1, nmo**2) tao = [] ao_loc = None moij, ijslice = _conc_mos(mo_coeff[i], mo_coeff[j])[2:] Lij_out = _ao2mo.r_e2(Lpq, moij, ijslice, tao, ao_loc, out=Lij_out) Lij.append(Lij_out.reshape(-1, nmo, nmo)) Lij = np.asarray(Lij) naux = Lij.shape[1] if kL == 0: km = kidx_r[kn] if len(idx[km]) > 0: for m in idx[km]: em = mo_energy[km][m] - omega # body dielectric matrix eps_body Pi = get_rho_response_R(gw, abs(em), mo_energy, Lij, kL, kidx) eps_body_inv = np.linalg.inv(np.eye(naux) - Pi) if gw.fc and m == orbp: # head dielectric matrix eps_00 Pi_00 = get_rho_response_head_R( gw, abs(em), mo_energy, qij) eps_00 = 1. - 4. * np.pi / np.linalg.norm( q_abs[0])**2 * Pi_00 # wings dielectric matrix eps_P0 Pi_P0 = get_rho_response_wing_R( gw, abs(em), mo_energy, Lij, qij) eps_P0 = -np.sqrt(4. * np.pi) / np.linalg.norm( q_abs[0]) * Pi_P0 # inverse dielectric matrix eps_inv_00 = 1. / (eps_00 - np.dot( np.dot(eps_P0.conj(), eps_body_inv), eps_P0)) eps_inv_P0 = -eps_inv_00 * np.dot( eps_body_inv, eps_P0) eps_inv_PQ = eps_body_inv # body Qmn = einsum('P,PQ->Q', Lij[km][:, m, orbp].conj(), eps_inv_PQ - np.eye(naux)) Wmn = 1. / nkpts * einsum('Q,Q->', Qmn, Lij[km][:, m, orbp]) sigmaR += fm * Wmn if gw.fc and m == orbp: # head correction Del_00 = 2. / np.pi * (6. * np.pi**2 / gw.mol.vol / nkpts)**(1. / 3.) * ( eps_inv_00 - 1.) sigmaR += fm * Del_00 # wings correction wings_const = np.sqrt( gw.mol.vol / 4. / np.pi**3) * (6. * np.pi**2 / gw.mol.vol / nkpts)**(2. / 3.) Wn_P0 = einsum('P,P->', Lij[kn][:, m, orbp].conj(), eps_inv_P0) Wn_P0 = Wn_P0.real * 2. sigmaR += fm * wings_const * Wn_P0 else: km = kidx_r[kn] if len(idx[km]) > 0: for m in idx[km]: em = mo_energy[km][m] - omega Pi = get_rho_response_R(gw, abs(em), mo_energy, Lij, kL, kidx) Pi_inv = np.linalg.inv(np.eye(naux) - Pi) - np.eye(naux) Qmn = einsum('P,PQ->Q', Lij[km][:, m, orbp].conj(), Pi_inv) Wmn = 1. / nkpts * einsum('Q,Q->', Qmn, Lij[km][:, m, orbp]) sigmaR += fm * Wmn return sigmaR
def get_WmnI_diag(gw, orbs, kptlist, freqs, max_memory=8000): ''' Compute GW correlation self-energy (diagonal elements) in MO basis on imaginary axis ''' mo_energy = np.array(gw._scf.mo_energy) mo_coeff = np.array(gw._scf.mo_coeff) nocc = gw.nocc nmo = gw.nmo nkpts = gw.nkpts kpts = gw.kpts nklist = len(kptlist) nw = len(freqs) norbs = len(orbs) mydf = gw.with_df # possible kpts shift center kscaled = gw.mol.get_scaled_kpts(kpts) kscaled -= kscaled[0] Del_00, Del_P0, qij, q_abs = None, None, None, None if gw.fc: # Set up q mesh for q->0 finite size correction q_pts = np.array([1e-3, 0, 0]).reshape(1, 3) nq_pts = len(q_pts) q_abs = gw.mol.get_abs_kpts(q_pts) # Get qij = 1/sqrt(Omega) * < psi_{ik} | e^{iqr} | psi_{ak-q} > at q: (nkpts, nocc, nvir) qij = get_qij(gw, q_abs[0], mo_coeff) Wmn = np.zeros((nkpts, nklist, nmo, norbs, nw), dtype=np.complex128) if gw.fc: Del_P0 = np.zeros((nklist, norbs, nw), dtype=np.complex128) Del_00 = np.zeros(nw, dtype=np.complex128) for kL in range(nkpts): # Lij: (ki, L, i, j) for looping every kL Lij = [] # kidx: save kj that conserves with kL and ki (-ki+kj+kL=G) # kidx_r: save ki that conserves with kL and kj (-ki+kj+kL=G) kidx = np.zeros((nkpts), dtype=np.int64) kidx_r = np.zeros((nkpts), dtype=np.int64) for i, kpti in enumerate(kpts): for j, kptj in enumerate(kpts): # Find (ki,kj) that satisfies momentum conservation with kL kconserv = -kscaled[i] + kscaled[j] + kscaled[kL] is_kconserv = np.linalg.norm(np.round(kconserv) - kconserv) < 1e-12 if is_kconserv: kidx[i] = j kidx_r[j] = i logger.debug( gw, "Read Lpq (kL: %s / %s, ki: %s, kj: %s)" % (kL + 1, nkpts, i, j)) Lij_out = None # Read (L|pq) and ao2mo transform to (L|ij) Lpq = [] for LpqR, LpqI, sign in mydf.sr_loop([kpti, kptj], max_memory=0.1 * gw._scf.max_memory, compact=False): Lpq.append(LpqR + LpqI * 1.0j) # support uneqaul naux on different k points Lpq = np.vstack(Lpq).reshape(-1, nmo**2) tao = [] ao_loc = None moij, ijslice = _conc_mos(mo_coeff[i], mo_coeff[j])[2:] Lij_out = _ao2mo.r_e2(Lpq, moij, ijslice, tao, ao_loc, out=Lij_out) Lij.append(Lij_out.reshape(-1, nmo, nmo)) Lij = np.asarray(Lij) naux = Lij.shape[1] if kL == 0: for w in range(nw): # body dielectric matrix eps_body Pi = get_rho_response(gw, freqs[w], mo_energy, Lij, kL, kidx) eps_body_inv = np.linalg.inv(np.eye(naux) - Pi) if gw.fc: # head dielectric matrix eps_00 Pi_00 = get_rho_response_head(gw, freqs[w], mo_energy, qij) eps_00 = 1. - 4. * np.pi / np.linalg.norm( q_abs[0])**2 * Pi_00 # wings dielectric matrix eps_P0 Pi_P0 = get_rho_response_wing(gw, freqs[w], mo_energy, Lij, qij) eps_P0 = -np.sqrt(4. * np.pi) / np.linalg.norm( q_abs[0]) * Pi_P0 # inverse dielectric matrix eps_inv_00 = 1. / (eps_00 - np.dot( np.dot(eps_P0.conj(), eps_body_inv), eps_P0)) eps_inv_P0 = -eps_inv_00 * np.dot(eps_body_inv, eps_P0) # head correction Del_00[w] = 2. / np.pi * (6. * np.pi**2 / gw.mol.vol / nkpts)**(1. / 3.) * (eps_inv_00 - 1.) wings_const = np.sqrt(gw.mol.vol / 4. / np.pi**3) * ( 6. * np.pi**2 / gw.mol.vol / nkpts)**(2. / 3.) eps_inv_PQ = eps_body_inv for k in range(nklist): kn = kptlist[k] # Find km that conserves with kn and kL (-km+kn+kL=G) km = kidx_r[kn] Qmn = einsum('Pmn,PQ->Qmn', Lij[km][:, :, orbs].conj(), eps_inv_PQ - np.eye(naux)) Wmn[km, k, :, :, w] = 1. / nkpts * einsum( 'Qmn,Qmn->mn', Qmn, Lij[km][:, :, orbs]) if gw.fc: # compute wing correction Wn_P0 = einsum('Pnm,P->nm', Lij[kn], eps_inv_P0).diagonal() Wn_P0 = Wn_P0.real * 2. Del_P0[k, :, w] = wings_const * Wn_P0[orbs] else: for w in range(nw): Pi = get_rho_response(gw, freqs[w], mo_energy, Lij, kL, kidx) Pi_inv = np.linalg.inv(np.eye(naux) - Pi) - np.eye(naux) for k in range(nklist): kn = kptlist[k] # Find km that conserves with kn and kL (-km+kn+kL=G) km = kidx_r[kn] Qmn = einsum('Pmn,PQ->Qmn', Lij[km][:, :, orbs].conj(), Pi_inv) Wmn[km, k, :, :, w] = 1. / nkpts * einsum( 'Qmn,Qmn->mn', Qmn, Lij[km][:, :, orbs]) return Wmn, Del_00, Del_P0, qij, q_abs
def _init_mp_df_eris(mp): """Compute 3-center electron repulsion integrals, i.e. (L|ov), where `L` denotes DF auxiliary basis functions and `o` and `v` occupied and virtual canonical crystalline orbitals. Note that `o` and `v` contain kpt indices `ko` and `kv`, and the third kpt index `kL` is determined by the conservation of momentum. Arguments: mp (KMP2) -- A KMP2 instance Returns: Lov (numpy.ndarray) -- 3-center DF ints, with shape (nkpts, nkpts, naux, nocc, nvir) """ from pyscf.pbc.df import df from pyscf.ao2mo import _ao2mo from pyscf.pbc.lib.kpts_helper import gamma_point log = logger.Logger(mp.stdout, mp.verbose) if mp._scf.with_df._cderi is None: mp._scf.with_df.build() cell = mp._scf.cell if cell.dimension == 2: # 2D ERIs are not positive definite. The 3-index tensors are stored in # two part. One corresponds to the positive part and one corresponds # to the negative part. The negative part is not considered in the # DF-driven CCSD implementation. raise NotImplementedError nocc = mp.nocc nmo = mp.nmo nvir = nmo - nocc nao = cell.nao_nr() mo_coeff = _add_padding(mp, mp.mo_coeff, mp.mo_energy)[0] kpts = mp.kpts nkpts = len(kpts) if gamma_point(kpts): dtype = np.double else: dtype = np.complex128 dtype = np.result_type(dtype, *mo_coeff) Lov = np.empty((nkpts, nkpts), dtype=object) cput0 = (logger.process_clock(), logger.perf_counter()) bra_start = 0 bra_end = nocc ket_start = nmo + nocc ket_end = ket_start + nvir with h5py.File(mp._scf.with_df._cderi, 'r') as f: kptij_lst = f['j3c-kptij'][:] tao = [] ao_loc = None for ki, kpti in enumerate(kpts): for kj, kptj in enumerate(kpts): kpti_kptj = np.array((kpti, kptj)) Lpq_ao = np.asarray(df._getitem(f, 'j3c', kpti_kptj, kptij_lst)) mo = np.hstack((mo_coeff[ki], mo_coeff[kj])) mo = np.asarray(mo, dtype=dtype, order='F') if dtype == np.double: out = _ao2mo.nr_e2( Lpq_ao, mo, (bra_start, bra_end, ket_start, ket_end), aosym='s2') else: #Note: Lpq.shape[0] != naux if linear dependency is found in auxbasis if Lpq_ao[0].size != nao**2: # aosym = 's2' Lpq_ao = lib.unpack_tril(Lpq_ao).astype(np.complex128) out = _ao2mo.r_e2(Lpq_ao, mo, (bra_start, bra_end, ket_start, ket_end), tao, ao_loc) Lov[ki, kj] = out.reshape(-1, nocc, nvir) log.timer_debug1("transforming DF-MP2 integrals", *cput0) return Lov
def general(mol, mo_coeffs, erifile, dataname='eri_mo', intor='int2e_spinor', aosym='s4', comp=None, max_memory=MAX_MEMORY, ioblk_size=IOBLK_SIZE, verbose=logger.WARN): time_0pass = (time.clock(), time.time()) log = logger.new_logger(mol, verbose) if '_spinor' not in intor: log.warn('r_ao2mo requires spinor integrals.\n' 'Suffix _spinor is added to %s', intor) intor = intor + '_spinor' intor, comp = gto.moleintor._get_intor_and_comp(mol._add_suffix(intor), comp) klsame = iden_coeffs(mo_coeffs[2], mo_coeffs[3]) nmoi = mo_coeffs[0].shape[1] nmoj = mo_coeffs[1].shape[1] nmok = mo_coeffs[2].shape[1] nmol = mo_coeffs[3].shape[1] nao = mo_coeffs[0].shape[0] aosym = outcore._stand_sym_code(aosym) if aosym in ('s1', 's2ij', 'a2ij'): nao_pair = nao * nao else: nao_pair = _count_naopair(mol, nao) nij_pair = nmoi*nmoj nkl_pair = nmok*nmol if klsame and aosym in ('s4', 's2kl', 'a2kl', 'a4ij', 'a4kl', 'a4'): log.debug('k-mo == l-mo') mokl = numpy.asarray(mo_coeffs[2], dtype=numpy.complex128, order='F') klshape = (0, nmok, 0, nmok) else: mokl = numpy.asarray(numpy.hstack((mo_coeffs[2],mo_coeffs[3])), dtype=numpy.complex128, order='F') klshape = (0, nmok, nmok, nmok+nmol) if isinstance(erifile, str): if h5py.is_hdf5(erifile): feri = h5py.File(erifile, 'a') if dataname in feri: del(feri[dataname]) else: feri = h5py.File(erifile, 'w') else: assert(isinstance(erifile, h5py.Group)) feri = erifile if comp == 1: chunks = (nmoj,nmol) shape = (nij_pair, nkl_pair) else: chunks = (1,nmoj,nmol) shape = (comp, nij_pair, nkl_pair) if nij_pair == 0 or nkl_pair == 0: feri.create_dataset(dataname, shape, 'c16') if isinstance(erifile, str): feri.close() return erifile else: h5d_eri = feri.create_dataset(dataname, shape, 'c16', chunks=chunks) log.debug('MO integrals %s are saved in %s/%s', intor, erifile, dataname) log.debug('num. MO ints = %.8g, required disk %.8g MB', float(nij_pair)*nkl_pair*comp, nij_pair*nkl_pair*comp*16/1e6) # transform e1 swapfile = tempfile.NamedTemporaryFile(dir=lib.param.TMPDIR) half_e1(mol, mo_coeffs, swapfile.name, intor, aosym, comp, max_memory, ioblk_size, log) time_1pass = log.timer('AO->MO transformation for %s 1 pass'%intor, *time_0pass) e2buflen = guess_e2bufsize(ioblk_size, nij_pair, nao_pair)[0] log.debug('step2: kl-pair (ao %d, mo %d), mem %.8g MB, ' 'ioblock (r/w) %.8g/%.8g MB', \ nao_pair, nkl_pair, e2buflen*nao_pair*16/1e6, e2buflen*nij_pair*16/1e6, e2buflen*nkl_pair*16/1e6) fswap = h5py.File(swapfile.name, 'r') klaoblks = len(fswap['0']) ijmoblks = int(numpy.ceil(float(nij_pair)/e2buflen)) * comp ao_loc = numpy.asarray(mol.ao_loc_2c(), dtype=numpy.int32) tao = numpy.asarray(mol.tmap(), dtype=numpy.int32) ti0 = time_1pass buf = numpy.empty((e2buflen, nao_pair), dtype=numpy.complex) istep = 0 for row0, row1 in prange(0, nij_pair, e2buflen): nrow = row1 - row0 for icomp in range(comp): istep += 1 tioi = 0 log.debug('step 2 [%d/%d], [%d,%d:%d], row = %d', \ istep, ijmoblks, icomp, row0, row1, nrow) col0 = 0 for ic in range(klaoblks): dat = fswap['%d/%d'%(icomp,ic)] col1 = col0 + dat.shape[1] buf[:nrow,col0:col1] = dat[row0:row1] col0 = col1 ti2 = log.timer('step 2 [%d/%d], load buf'%(istep,ijmoblks), *ti0) tioi += ti2[1]-ti0[1] pbuf = _ao2mo.r_e2(buf[:nrow], mokl, klshape, tao, ao_loc, aosym) tw1 = time.time() if comp == 1: h5d_eri[row0:row1] = pbuf else: h5d_eri[icomp,row0:row1] = pbuf tioi += time.time()-tw1 ti1 = (time.clock(), time.time()) log.debug('step 2 [%d/%d] CPU time: %9.2f, Wall time: %9.2f, I/O time: %9.2f', \ istep, ijmoblks, ti1[0]-ti0[0], ti1[1]-ti0[1], tioi) ti0 = ti1 buf = pbuf = None fswap.close() if isinstance(erifile, str): feri.close() log.timer('AO->MO transformation for %s 2 pass'%intor, *time_1pass) log.timer('AO->MO transformation for %s '%intor, *time_0pass) return erifile
def ao2mo_7d(mydf, mo_coeff_kpts, kpts=None, factor=1, out=None): cell = mydf.cell if kpts is None: kpts = mydf.kpts nkpts = len(kpts) if isinstance(mo_coeff_kpts, numpy.ndarray) and mo_coeff_kpts.ndim == 3: mo_coeff_kpts = [mo_coeff_kpts] * 4 else: mo_coeff_kpts = list(mo_coeff_kpts) # Shape of the orbitals can be different on different k-points. The # orbital coefficients must be formatted (padded by zeros) so that the # shape of the orbital coefficients are the same on all k-points. This can # be achieved by calling pbc.mp.kmp2.padded_mo_coeff function nmoi, nmoj, nmok, nmol = [x.shape[2] for x in mo_coeff_kpts] eri_shape = (nkpts, nkpts, nkpts, nmoi, nmoj, nmok, nmol) if gamma_point(kpts): dtype = numpy.result_type(*mo_coeff_kpts) else: dtype = numpy.complex128 if out is None: out = numpy.empty(eri_shape, dtype=dtype) else: assert (out.shape == eri_shape) kptij_lst = numpy.array([(ki, kj) for ki in kpts for kj in kpts]) kptis_lst = kptij_lst[:, 0] kptjs_lst = kptij_lst[:, 1] kpt_ji = kptjs_lst - kptis_lst uniq_kpts, uniq_index, uniq_inverse = unique(kpt_ji) ngrids = numpy.prod(mydf.mesh) nao = cell.nao_nr() max_memory = max( 2000, mydf.max_memory - lib.current_memory()[0] - nao**4 * 16 / 1e6) * .5 fswap = lib.H5TmpFile() tao = [] ao_loc = None kconserv = kpts_helper.get_kconserv(cell, kpts) for uniq_id, kpt in enumerate(uniq_kpts): q = uniq_kpts[uniq_id] adapted_ji_idx = numpy.where(uniq_inverse == uniq_id)[0] kptjs = kptjs_lst[adapted_ji_idx] coulG = mydf.weighted_coulG(q, False, mydf.mesh) coulG *= factor moij_list = [] ijslice_list = [] for ji, ji_idx in enumerate(adapted_ji_idx): ki = ji_idx // nkpts kj = ji_idx % nkpts moij, ijslice = _conc_mos(mo_coeff_kpts[0][ki], mo_coeff_kpts[1][kj])[2:] moij_list.append(moij) ijslice_list.append(ijslice) fswap.create_dataset('zij/' + str(ji), (ngrids, nmoi * nmoj), 'D') for aoaoks, p0, p1 in mydf.ft_loop(mydf.mesh, q, kptjs, max_memory=max_memory): for ji, aoao in enumerate(aoaoks): ki = adapted_ji_idx[ji] // nkpts kj = adapted_ji_idx[ji] % nkpts buf = aoao.transpose(1, 2, 0).reshape(nao**2, p1 - p0) zij = _ao2mo.r_e2(lib.transpose(buf), moij_list[ji], ijslice_list[ji], tao, ao_loc) zij *= coulG[p0:p1, None] fswap['zij/' + str(ji)][p0:p1] = zij mokl_list = [] klslice_list = [] for kk in range(nkpts): kl = kconserv[ki, kj, kk] mokl, klslice = _conc_mos(mo_coeff_kpts[2][kk], mo_coeff_kpts[3][kl])[2:] mokl_list.append(mokl) klslice_list.append(klslice) fswap.create_dataset('zkl/' + str(kk), (ngrids, nmok * nmol), 'D') ki = adapted_ji_idx[0] // nkpts kj = adapted_ji_idx[0] % nkpts kptls = kpts[kconserv[ki, kj, :]] for aoaoks, p0, p1 in mydf.ft_loop(mydf.mesh, q, -kptls, max_memory=max_memory): for kk, aoao in enumerate(aoaoks): buf = aoao.conj().transpose(1, 2, 0).reshape(nao**2, p1 - p0) zkl = _ao2mo.r_e2(lib.transpose(buf), mokl_list[kk], klslice_list[kk], tao, ao_loc) fswap['zkl/' + str(kk)][p0:p1] = zkl for ji, ji_idx in enumerate(adapted_ji_idx): ki = ji_idx // nkpts kj = ji_idx % nkpts moij, ijslice = _conc_mos(mo_coeff_kpts[0][ki], mo_coeff_kpts[1][kj])[2:] zij = [] for LpqR, LpqI, sign in mydf.sr_loop(kpts[[ki, kj]], max_memory, False, mydf.blockdim): zij.append( _ao2mo.r_e2(LpqR + LpqI * 1j, moij, ijslice, tao, ao_loc)) for kk in range(nkpts): kl = kconserv[ki, kj, kk] eri_mo = lib.dot( numpy.asarray(fswap['zij/' + str(ji)]).T, numpy.asarray(fswap['zkl/' + str(kk)])) for i, (LrsR, LrsI, sign) in \ enumerate(mydf.sr_loop(kpts[[kk,kl]], max_memory, False, mydf.blockdim)): zkl = _ao2mo.r_e2(LrsR + LrsI * 1j, mokl_list[kk], klslice_list[kk], tao, ao_loc) lib.dot(zij[i].T, zkl, sign * factor, eri_mo, 1) if dtype == numpy.double: eri_mo = eri_mo.real out[ki, kj, kk] = eri_mo.reshape(eri_shape[3:]) del (fswap['zij']) del (fswap['zkl']) return out
def general(mydf, mo_coeffs, kpts=None, compact=getattr(__config__, 'pbc_df_ao2mo_general_compact', True)): warn_pbc2d_eri(mydf) cell = mydf.cell kptijkl = _format_kpts(kpts) kpti, kptj, kptk, kptl = kptijkl if isinstance(mo_coeffs, numpy.ndarray) and mo_coeffs.ndim == 2: mo_coeffs = (mo_coeffs,) * 4 if not _iskconserv(cell, kptijkl): lib.logger.warn(cell, 'aft_ao2mo: momentum conservation not found in ' 'the given k-points %s', kptijkl) return numpy.zeros([mo.shape[1] for mo in mo_coeffs]) q = kptj - kpti mesh = mydf.mesh coulG = mydf.weighted_coulG(q, False, mesh) all_real = not any(numpy.iscomplexobj(mo) for mo in mo_coeffs) max_memory = max(2000, (mydf.max_memory - lib.current_memory()[0]) * .5) #################### # gamma point, the integral is real and with s4 symmetry if gamma_point(kptijkl) and all_real: ijmosym, nij_pair, moij, ijslice = _conc_mos(mo_coeffs[0], mo_coeffs[1], compact) klmosym, nkl_pair, mokl, klslice = _conc_mos(mo_coeffs[2], mo_coeffs[3], compact) eri_mo = numpy.zeros((nij_pair,nkl_pair)) sym = (iden_coeffs(mo_coeffs[0], mo_coeffs[2]) and iden_coeffs(mo_coeffs[1], mo_coeffs[3])) ijR = ijI = klR = klI = buf = None for pqkR, pqkI, p0, p1 \ in mydf.pw_loop(mesh, kptijkl[:2], q, max_memory=max_memory, aosym='s2'): buf = lib.transpose(pqkR, out=buf) ijR, klR = _dtrans(buf, ijR, ijmosym, moij, ijslice, buf, klR, klmosym, mokl, klslice, sym) lib.ddot(ijR.T, klR*coulG[p0:p1,None], 1, eri_mo, 1) buf = lib.transpose(pqkI, out=buf) ijI, klI = _dtrans(buf, ijI, ijmosym, moij, ijslice, buf, klI, klmosym, mokl, klslice, sym) lib.ddot(ijI.T, klI*coulG[p0:p1,None], 1, eri_mo, 1) pqkR = pqkI = None return eri_mo #################### # (kpt) i == j == k == l != 0 # (kpt) i == l && j == k && i != j && j != k => # elif is_zero(kpti-kptl) and is_zero(kptj-kptk): mo_coeffs = _mo_as_complex(mo_coeffs) nij_pair, moij, ijslice = _conc_mos(mo_coeffs[0], mo_coeffs[1])[1:] nlk_pair, molk, lkslice = _conc_mos(mo_coeffs[3], mo_coeffs[2])[1:] eri_mo = numpy.zeros((nij_pair,nlk_pair), dtype=numpy.complex) sym = (iden_coeffs(mo_coeffs[0], mo_coeffs[3]) and iden_coeffs(mo_coeffs[1], mo_coeffs[2])) zij = zlk = buf = None for pqkR, pqkI, p0, p1 \ in mydf.pw_loop(mesh, kptijkl[:2], q, max_memory=max_memory): buf = lib.transpose(pqkR+pqkI*1j, out=buf) zij, zlk = _ztrans(buf, zij, moij, ijslice, buf, zlk, molk, lkslice, sym) lib.dot(zij.T, zlk.conj()*coulG[p0:p1,None], 1, eri_mo, 1) pqkR = pqkI = None nmok = mo_coeffs[2].shape[1] nmol = mo_coeffs[3].shape[1] eri_mo = lib.transpose(eri_mo.reshape(-1,nmol,nmok), axes=(0,2,1)) return eri_mo.reshape(nij_pair,nlk_pair) #################### # aosym = s1, complex integrals # # If kpti == kptj, (kptl-kptk)*a has to be multiples of 2pi because of the wave # vector symmetry. k is a fraction of reciprocal basis, 0 < k/b < 1, by definition. # So kptl/b - kptk/b must be -1 < k/b < 1. => kptl == kptk # else: mo_coeffs = _mo_as_complex(mo_coeffs) nij_pair, moij, ijslice = _conc_mos(mo_coeffs[0], mo_coeffs[1])[1:] nkl_pair, mokl, klslice = _conc_mos(mo_coeffs[2], mo_coeffs[3])[1:] eri_mo = numpy.zeros((nij_pair,nkl_pair), dtype=numpy.complex) tao = [] ao_loc = None zij = zkl = buf = None for (pqkR, pqkI, p0, p1), (rskR, rskI, q0, q1) in \ lib.izip(mydf.pw_loop(mesh, kptijkl[:2], q, max_memory=max_memory*.5), mydf.pw_loop(mesh,-kptijkl[2:], q, max_memory=max_memory*.5)): buf = lib.transpose(pqkR+pqkI*1j, out=buf) zij = _ao2mo.r_e2(buf, moij, ijslice, tao, ao_loc, out=zij) buf = lib.transpose(rskR-rskI*1j, out=buf) zkl = _ao2mo.r_e2(buf, mokl, klslice, tao, ao_loc, out=zkl) zij *= coulG[p0:p1,None] lib.dot(zij.T, zkl, 1, eri_mo, 1) pqkR = pqkI = rskR = rskI = None return eri_mo
def _init_cis_df_eris(cis, eris): """Add 3-center electron repulsion integrals, i.e. (L|pq), in `eris`, where `L` denotes DF auxiliary basis functions and `p` and `q` canonical crystalline orbitals. Note that `p` and `q` contain kpt indices `kp` and `kq`, and the third kpt index `kL` is determined by the conservation of momentum. Arguments: cis {KCIS} -- A KCIS instance eris {_CIS_ERIS} -- A _CIS_ERIS instance to which we want to add 3c ints Returns: _CIS_ERIS -- A _CIS_ERIS instance with 3c ints """ from pyscf.pbc.df import df from pyscf.ao2mo import _ao2mo from pyscf.pbc.lib.kpts_helper import gamma_point log = logger.Logger(cis.stdout, cis.verbose) if cis._scf.with_df._cderi is None: cis._scf.with_df.build() cell = cis._scf.cell if cell.dimension == 2: # 2D ERIs are not positive definite. The 3-index tensors are stored in # two part. One corresponds to the positive part and one corresponds # to the negative part. The negative part is not considered in the # DF-driven CCSD implementation. raise NotImplementedError nocc = cis.nocc nmo = cis.nmo nao = cell.nao_nr() kpts = cis.kpts nkpts = len(kpts) if gamma_point(kpts): dtype = np.double else: dtype = np.complex128 eris.dtype = dtype = np.result_type(dtype, *eris.mo_coeff) eris.Lpq_mo = Lpq_mo = np.empty((nkpts, nkpts), dtype=object) cput0 = (time.clock(), time.time()) with h5py.File(cis._scf.with_df._cderi, 'r') as f: kptij_lst = f['j3c-kptij'].value tao = [] ao_loc = None for ki, kpti in enumerate(kpts): for kj, kptj in enumerate(kpts): kpti_kptj = np.array((kpti, kptj)) Lpq_ao = np.asarray(df._getitem(f, 'j3c', kpti_kptj, kptij_lst)) mo = np.hstack((eris.mo_coeff[ki], eris.mo_coeff[kj])) mo = np.asarray(mo, dtype=dtype, order='F') if dtype == np.double: out = _ao2mo.nr_e2(Lpq_ao, mo, (0, nmo, nmo, nmo + nmo), aosym='s2') else: #Note: Lpq.shape[0] != naux if linear dependency is found in auxbasis if Lpq_ao[0].size != nao**2: # aosym = 's2' Lpq_ao = lib.unpack_tril(Lpq_ao).astype(np.complex128) out = _ao2mo.r_e2(Lpq_ao, mo, (0, nmo, nmo, nmo + nmo), tao, ao_loc) Lpq_mo[ki, kj] = out.reshape(-1, nmo, nmo) log.timer_debug1("transforming DF-CIS integrals", *cput0) return eris
def ao2mo_7d(mydf, mo_coeff_kpts, kpts=None, factor=1, out=None): cell = mydf.cell if kpts is None: kpts = mydf.kpts nkpts = len(kpts) if isinstance(mo_coeff_kpts, numpy.ndarray) and mo_coeff_kpts.ndim == 3: mo_coeff_kpts = [mo_coeff_kpts] * 4 else: mo_coeff_kpts = list(mo_coeff_kpts) # Shape of the orbitals can be different on different k-points. The # orbital coefficients must be formatted (padded by zeros) so that the # shape of the orbital coefficients are the same on all k-points. This can # be achieved by calling pbc.mp.kmp2.padded_mo_coeff function nmoi, nmoj, nmok, nmol = [x.shape[2] for x in mo_coeff_kpts] eri_shape = (nkpts, nkpts, nkpts, nmoi, nmoj, nmok, nmol) if gamma_point(kpts): dtype = numpy.result_type(*mo_coeff_kpts) else: dtype = numpy.complex128 if out is None: out = numpy.empty(eri_shape, dtype=dtype) else: assert(out.shape == eri_shape) kptij_lst = numpy.array([(ki, kj) for ki in kpts for kj in kpts]) kptis_lst = kptij_lst[:,0] kptjs_lst = kptij_lst[:,1] kpt_ji = kptjs_lst - kptis_lst uniq_kpts, uniq_index, uniq_inverse = unique(kpt_ji) ngrids = numpy.prod(mydf.mesh) nao = cell.nao_nr() max_memory = max(2000, mydf.max_memory-lib.current_memory()[0]-nao**4*16/1e6) * .5 fswap = lib.H5TmpFile() tao = [] ao_loc = None kconserv = kpts_helper.get_kconserv(cell, kpts) for uniq_id, kpt in enumerate(uniq_kpts): q = uniq_kpts[uniq_id] adapted_ji_idx = numpy.where(uniq_inverse == uniq_id)[0] kptjs = kptjs_lst[adapted_ji_idx] coulG = mydf.weighted_coulG(q, False, mydf.mesh) coulG *= factor moij_list = [] ijslice_list = [] for ji, ji_idx in enumerate(adapted_ji_idx): ki = ji_idx // nkpts kj = ji_idx % nkpts moij, ijslice = _conc_mos(mo_coeff_kpts[0][ki], mo_coeff_kpts[1][kj])[2:] moij_list.append(moij) ijslice_list.append(ijslice) fswap.create_dataset('zij/'+str(ji), (ngrids,nmoi*nmoj), 'D') for aoaoks, p0, p1 in mydf.ft_loop(mydf.mesh, q, kptjs): for ji, aoao in enumerate(aoaoks): ki = adapted_ji_idx[ji] // nkpts kj = adapted_ji_idx[ji] % nkpts buf = aoao.transpose(1,2,0).reshape(nao**2,ngrids) zij = _ao2mo.r_e2(lib.transpose(buf), moij_list[ji], ijslice_list[ji], tao, ao_loc) zij *= coulG[p0:p1,None] fswap['zij/'+str(ji)][p0:p1] = zij mokl_list = [] klslice_list = [] for kk in range(nkpts): kl = kconserv[ki, kj, kk] mokl, klslice = _conc_mos(mo_coeff_kpts[2][kk], mo_coeff_kpts[3][kl])[2:] mokl_list.append(mokl) klslice_list.append(klslice) fswap.create_dataset('zkl/'+str(kk), (ngrids,nmok*nmol), 'D') ki = adapted_ji_idx[0] // nkpts kj = adapted_ji_idx[0] % nkpts kptls = kpts[kconserv[ki, kj, :]] for aoaoks, p0, p1 in mydf.ft_loop(mydf.mesh, q, -kptls): for kk, aoao in enumerate(aoaoks): buf = aoao.conj().transpose(1,2,0).reshape(nao**2,ngrids) zkl = _ao2mo.r_e2(lib.transpose(buf), mokl_list[kk], klslice_list[kk], tao, ao_loc) fswap['zkl/'+str(kk)][p0:p1] = zkl for ji, ji_idx in enumerate(adapted_ji_idx): ki = ji_idx // nkpts kj = ji_idx % nkpts moij, ijslice = _conc_mos(mo_coeff_kpts[0][ki], mo_coeff_kpts[1][kj])[2:] zij = [] for LpqR, LpqI, sign in mydf.sr_loop(kpts[[ki,kj]], max_memory, False, mydf.blockdim): zij.append(_ao2mo.r_e2(LpqR+LpqI*1j, moij, ijslice, tao, ao_loc)) for kk in range(nkpts): kl = kconserv[ki, kj, kk] eri_mo = lib.dot(numpy.asarray(fswap['zij/'+str(ji)]).T, numpy.asarray(fswap['zkl/'+str(kk)])) for i, (LrsR, LrsI, sign) in \ enumerate(mydf.sr_loop(kpts[[kk,kl]], max_memory, False, mydf.blockdim)): zkl = _ao2mo.r_e2(LrsR+LrsI*1j, mokl_list[kk], klslice_list[kk], tao, ao_loc) lib.dot(zij[i].T, zkl, sign*factor, eri_mo, 1) if dtype == numpy.double: eri_mo = eri_mo.real out[ki,kj,kk] = eri_mo.reshape(eri_shape[3:]) del(fswap['zij']) del(fswap['zkl']) return out
def ao2mo_7d(mydf, mo_coeff_kpts, kpts=None, factor=1, out=None): cell = mydf.cell if kpts is None: kpts = mydf.kpts nkpts = len(kpts) if isinstance(mo_coeff_kpts, numpy.ndarray) and mo_coeff_kpts.ndim == 3: mo_coeff_kpts = [mo_coeff_kpts] * 4 else: mo_coeff_kpts = list(mo_coeff_kpts) # Shape of the orbitals can be different on different k-points. The # orbital coefficients must be formatted (padded by zeros) so that the # shape of the orbital coefficients are the same on all k-points. This can # be achieved by calling pbc.mp.kmp2.padded_mo_coeff function nmoi, nmoj, nmok, nmol = [x.shape[2] for x in mo_coeff_kpts] eri_shape = (nkpts, nkpts, nkpts, nmoi, nmoj, nmok, nmol) if gamma_point(kpts): dtype = numpy.result_type(*mo_coeff_kpts) else: dtype = numpy.complex128 if out is None: out = numpy.empty(eri_shape, dtype=dtype) else: assert(out.shape == eri_shape) kptij_lst = numpy.array([(ki, kj) for ki in kpts for kj in kpts]) kptis_lst = kptij_lst[:,0] kptjs_lst = kptij_lst[:,1] kpt_ji = kptjs_lst - kptis_lst uniq_kpts, uniq_index, uniq_inverse = unique(kpt_ji) ngrids = numpy.prod(mydf.mesh) nao = cell.nao_nr() max_memory = max(2000, mydf.max_memory-lib.current_memory()[0]-nao**4*16/1e6) * .5 tao = [] ao_loc = None kconserv = kpts_helper.get_kconserv(cell, kpts) for uniq_id, kpt in enumerate(uniq_kpts): q = uniq_kpts[uniq_id] adapted_ji_idx = numpy.where(uniq_inverse == uniq_id)[0] for ji, ji_idx in enumerate(adapted_ji_idx): ki = ji_idx // nkpts kj = ji_idx % nkpts moij, ijslice = _conc_mos(mo_coeff_kpts[0][ki], mo_coeff_kpts[1][kj])[2:] zij = [] for LpqR, LpqI, sign in mydf.sr_loop(kpts[[ki,kj]], max_memory, False, mydf.blockdim): zij.append(_ao2mo.r_e2(LpqR+LpqI*1j, moij, ijslice, tao, ao_loc)) for kk in range(nkpts): kl = kconserv[ki, kj, kk] mokl, klslice = _conc_mos(mo_coeff_kpts[2][kk], mo_coeff_kpts[3][kl])[2:] eri_mo = numpy.zeros((nmoi*nmoj,nmok*nmol), dtype=numpy.complex128) for i, (LrsR, LrsI, sign) in \ enumerate(mydf.sr_loop(kpts[[kk,kl]], max_memory, False, mydf.blockdim)): zkl = _ao2mo.r_e2(LrsR+LrsI*1j, mokl, klslice, tao, ao_loc) lib.dot(zij[i].T, zkl, sign*factor, eri_mo, 1) if dtype == numpy.double: eri_mo = eri_mo.real out[ki,kj,kk] = eri_mo.reshape(eri_shape[3:]) return out
def ao2mo_7d(mydf, mo_coeff_kpts, kpts=None, factor=1, out=None): cell = mydf.cell if kpts is None: kpts = mydf.kpts nkpts = len(kpts) if isinstance(mo_coeff_kpts, numpy.ndarray) and mo_coeff_kpts.ndim == 3: mo_coeff_kpts = [mo_coeff_kpts] * 4 else: mo_coeff_kpts = list(mo_coeff_kpts) # Shape of the orbitals can be different on different k-points. The # orbital coefficients must be formatted (padded by zeros) so that the # shape of the orbital coefficients are the same on all k-points. This can # be achieved by calling pbc.mp.kmp2.padded_mo_coeff function nmoi, nmoj, nmok, nmol = [x.shape[2] for x in mo_coeff_kpts] eri_shape = (nkpts, nkpts, nkpts, nmoi, nmoj, nmok, nmol) if gamma_point(kpts): dtype = numpy.result_type(*mo_coeff_kpts) else: dtype = numpy.complex128 if out is None: out = numpy.empty(eri_shape, dtype=dtype) else: assert(out.shape == eri_shape) kptij_lst = numpy.array([(ki, kj) for ki in kpts for kj in kpts]) kptis_lst = kptij_lst[:,0] kptjs_lst = kptij_lst[:,1] kpt_ji = kptjs_lst - kptis_lst uniq_kpts, uniq_index, uniq_inverse = unique(kpt_ji) nao = cell.nao_nr() max_memory = max(2000, mydf.max_memory-lib.current_memory()[0]-nao**4*16/1e6) * .5 tao = [] ao_loc = None kconserv = kpts_helper.get_kconserv(cell, kpts) for uniq_id, kpt in enumerate(uniq_kpts): adapted_ji_idx = numpy.where(uniq_inverse == uniq_id)[0] for ji, ji_idx in enumerate(adapted_ji_idx): ki = ji_idx // nkpts kj = ji_idx % nkpts moij, ijslice = _conc_mos(mo_coeff_kpts[0][ki], mo_coeff_kpts[1][kj])[2:] zij = [] for LpqR, LpqI, sign in mydf.sr_loop(kpts[[ki,kj]], max_memory, False, mydf.blockdim): zij.append(_ao2mo.r_e2(LpqR+LpqI*1j, moij, ijslice, tao, ao_loc)) for kk in range(nkpts): kl = kconserv[ki, kj, kk] mokl, klslice = _conc_mos(mo_coeff_kpts[2][kk], mo_coeff_kpts[3][kl])[2:] eri_mo = numpy.zeros((nmoi*nmoj,nmok*nmol), dtype=numpy.complex128) for i, (LrsR, LrsI, sign) in \ enumerate(mydf.sr_loop(kpts[[kk,kl]], max_memory, False, mydf.blockdim)): zkl = _ao2mo.r_e2(LrsR+LrsI*1j, mokl, klslice, tao, ao_loc) lib.dot(zij[i].T, zkl, sign*factor, eri_mo, 1) if dtype == numpy.double: eri_mo = eri_mo.real out[ki,kj,kk] = eri_mo.reshape(eri_shape[3:]) return out
def general(mol, mo_coeffs, erifile, dataname='eri_mo', tmpdir=None, intor='cint2e', aosym='s4', comp=1, max_memory=4000, ioblk_size=256, verbose=logger.WARN): time_0pass = (time.clock(), time.time()) if isinstance(verbose, logger.Logger): log = verbose else: log = logger.Logger(mol.stdout, verbose) klsame = iden_coeffs(mo_coeffs[2], mo_coeffs[3]) nmoi = mo_coeffs[0].shape[1] nmoj = mo_coeffs[1].shape[1] nmok = mo_coeffs[2].shape[1] nmol = mo_coeffs[3].shape[1] nao = mo_coeffs[0].shape[0] aosym = _stand_sym_code(aosym) if aosym in ('s1', 's2ij', 'a2ij'): nao_pair = nao * nao else: nao_pair = guess_nao_pair(mol, nao) nij_pair = nmoi*nmoj nkl_pair = nmok*nmol if klsame and aosym in ('s4', 's2kl', 'a2kl', 'a4ij', 'a4kl', 'a4'): log.debug('k-mo == l-mo') mokl = numpy.asarray(mo_coeffs[2], order='F') klshape = (0, nmok, 0, nmok) else: mokl = numpy.asarray(numpy.hstack((mo_coeffs[2],mo_coeffs[3])), order='F') klshape = (0, nmok, nmok, nmok+nmol) if isinstance(erifile, str): if h5py.is_hdf5(erifile): feri = h5py.File(erifile) if dataname in feri: del(feri[dataname]) else: feri = h5py.File(erifile, 'w') else: assert(isinstance(erifile, h5py.Group)) feri = erifile if comp == 1: chunks = (nmoj,nmol) h5d_eri = feri.create_dataset(dataname, (nij_pair,nkl_pair), 'c16', chunks=chunks) else: chunks = (1,nmoj,nmol) h5d_eri = feri.create_dataset(dataname, (comp,nij_pair,nkl_pair), 'c16', chunks=chunks) if nij_pair == 0 or nkl_pair == 0: feri.close() return erifile log.debug('MO integrals %s are saved in %s/%s', intor, erifile, dataname) log.debug('num. MO ints = %.8g, required disk %.8g MB', float(nij_pair)*nkl_pair*comp, nij_pair*nkl_pair*comp*16/1e6) # transform e1 swapfile = tempfile.NamedTemporaryFile(dir=tmpdir) half_e1(mol, mo_coeffs, swapfile.name, intor, aosym, comp, max_memory, ioblk_size, log) time_1pass = log.timer('AO->MO transformation for %s 1 pass'%intor, *time_0pass) e2buflen = guess_e2bufsize(ioblk_size, nij_pair, nao_pair)[0] log.debug('step2: kl-pair (ao %d, mo %d), mem %.8g MB, ' 'ioblock (r/w) %.8g/%.8g MB', \ nao_pair, nkl_pair, e2buflen*nao_pair*16/1e6, e2buflen*nij_pair*16/1e6, e2buflen*nkl_pair*16/1e6) fswap = h5py.File(swapfile.name, 'r') klaoblks = len(fswap['0']) ijmoblks = int(numpy.ceil(float(nij_pair)/e2buflen)) * comp ao_loc = numpy.asarray(mol.ao_loc_2c(), dtype=numpy.int32) tao = numpy.asarray(mol.tmap(), dtype=numpy.int32) ti0 = time_1pass buf = numpy.empty((e2buflen, nao_pair), dtype=numpy.complex) istep = 0 for row0, row1 in prange(0, nij_pair, e2buflen): nrow = row1 - row0 for icomp in range(comp): istep += 1 tioi = 0 log.debug('step 2 [%d/%d], [%d,%d:%d], row = %d', \ istep, ijmoblks, icomp, row0, row1, nrow) col0 = 0 for ic in range(klaoblks): dat = fswap['%d/%d'%(icomp,ic)] col1 = col0 + dat.shape[1] buf[:nrow,col0:col1] = dat[row0:row1] col0 = col1 ti2 = log.timer('step 2 [%d/%d], load buf'%(istep,ijmoblks), *ti0) tioi += ti2[1]-ti0[1] pbuf = _ao2mo.r_e2(buf[:nrow], mokl, klshape, tao, ao_loc, aosym) tw1 = time.time() if comp == 1: h5d_eri[row0:row1] = pbuf else: h5d_eri[icomp,row0:row1] = pbuf tioi += time.time()-tw1 ti1 = (time.clock(), time.time()) log.debug('step 2 [%d/%d] CPU time: %9.2f, Wall time: %9.2f, I/O time: %9.2f', \ istep, ijmoblks, ti1[0]-ti0[0], ti1[1]-ti0[1], tioi) ti0 = ti1 buf = pbuf = None fswap.close() if isinstance(erifile, str): feri.close() log.timer('AO->MO transformation for %s 2 pass'%intor, *time_1pass) log.timer('AO->MO transformation for %s '%intor, *time_0pass) return erifile
def general(mydf, mo_coeffs, kpts=None, compact=True): kptijkl = _format_kpts(kpts) kpti, kptj, kptk, kptl = kptijkl if isinstance(mo_coeffs, numpy.ndarray) and mo_coeffs.ndim == 2: mo_coeffs = (mo_coeffs, ) * 4 q = kptj - kpti coulG = mydf.weighted_coulG(q, False, mydf.gs) all_real = not any(numpy.iscomplexobj(mo) for mo in mo_coeffs) max_memory = max(2000, (mydf.max_memory - lib.current_memory()[0]) * .5) #################### # gamma point, the integral is real and with s4 symmetry if gamma_point(kptijkl) and all_real: ijmosym, nij_pair, moij, ijslice = _conc_mos(mo_coeffs[0], mo_coeffs[1], compact) klmosym, nkl_pair, mokl, klslice = _conc_mos(mo_coeffs[2], mo_coeffs[3], compact) eri_mo = numpy.zeros((nij_pair, nkl_pair)) sym = (iden_coeffs(mo_coeffs[0], mo_coeffs[2]) and iden_coeffs(mo_coeffs[1], mo_coeffs[3])) ijR = ijI = klR = klI = buf = None for pqkR, pqkI, p0, p1 \ in mydf.pw_loop(mydf.gs, kptijkl[:2], q, max_memory=max_memory, aosym='s2'): vG = numpy.sqrt(coulG[p0:p1]) pqkR *= vG pqkI *= vG buf = lib.transpose(pqkR, out=buf) ijR, klR = _dtrans(buf, ijR, ijmosym, moij, ijslice, buf, klR, klmosym, mokl, klslice, sym) lib.ddot(ijR.T, klR, 1, eri_mo, 1) buf = lib.transpose(pqkI, out=buf) ijI, klI = _dtrans(buf, ijI, ijmosym, moij, ijslice, buf, klI, klmosym, mokl, klslice, sym) lib.ddot(ijI.T, klI, 1, eri_mo, 1) pqkR = pqkI = None return eri_mo #################### # (kpt) i == j == k == l != 0 # (kpt) i == l && j == k && i != j && j != k => # elif is_zero(kpti - kptl) and is_zero(kptj - kptk): mo_coeffs = _mo_as_complex(mo_coeffs) nij_pair, moij, ijslice = _conc_mos(mo_coeffs[0], mo_coeffs[1])[1:] nlk_pair, molk, lkslice = _conc_mos(mo_coeffs[3], mo_coeffs[2])[1:] eri_mo = numpy.zeros((nij_pair, nlk_pair), dtype=numpy.complex) sym = (iden_coeffs(mo_coeffs[0], mo_coeffs[3]) and iden_coeffs(mo_coeffs[1], mo_coeffs[2])) zij = zlk = buf = None for pqkR, pqkI, p0, p1 \ in mydf.pw_loop(mydf.gs, kptijkl[:2], q, max_memory=max_memory): buf = lib.transpose(pqkR + pqkI * 1j, out=buf) buf *= numpy.sqrt(coulG[p0:p1]).reshape(-1, 1) zij, zlk = _ztrans(buf, zij, moij, ijslice, buf, zlk, molk, lkslice, sym) lib.dot(zij.T, zlk.conj(), 1, eri_mo, 1) pqkR = pqkI = None nmok = mo_coeffs[2].shape[1] nmol = mo_coeffs[3].shape[1] eri_mo = lib.transpose(eri_mo.reshape(-1, nmol, nmok), axes=(0, 2, 1)) return eri_mo.reshape(nij_pair, nlk_pair) #################### # aosym = s1, complex integrals # # If kpti == kptj, (kptl-kptk)*a has to be multiples of 2pi because of the wave # vector symmetry. k is a fraction of reciprocal basis, 0 < k/b < 1, by definition. # So kptl/b - kptk/b must be -1 < k/b < 1. => kptl == kptk # else: mo_coeffs = _mo_as_complex(mo_coeffs) nij_pair, moij, ijslice = _conc_mos(mo_coeffs[0], mo_coeffs[1])[1:] nkl_pair, mokl, klslice = _conc_mos(mo_coeffs[2], mo_coeffs[3])[1:] eri_mo = numpy.zeros((nij_pair, nkl_pair), dtype=numpy.complex) tao = [] ao_loc = None zij = zkl = buf = None for (pqkR, pqkI, p0, p1), (rskR, rskI, q0, q1) in \ lib.izip(mydf.pw_loop(mydf.gs, kptijkl[:2], q, max_memory=max_memory*.5), mydf.pw_loop(mydf.gs,-kptijkl[2:], q, max_memory=max_memory*.5)): buf = lib.transpose(pqkR + pqkI * 1j, out=buf) zij = _ao2mo.r_e2(buf, moij, ijslice, tao, ao_loc, out=zij) buf = lib.transpose(rskR - rskI * 1j, out=buf) zkl = _ao2mo.r_e2(buf, mokl, klslice, tao, ao_loc, out=zkl) zij *= coulG[p0:p1].reshape(-1, 1) lib.dot(zij.T, zkl, 1, eri_mo, 1) pqkR = pqkI = rskR = rskI = None return eri_mo
def get_sigma_diag(gw, orbs, kptlist, freqs, wts, iw_cutoff=None, max_memory=8000): ''' Compute GW correlation self-energy (diagonal elements) in MO basis on imaginary axis ''' mo_energy = np.array(gw._scf.mo_energy) mo_coeff = np.array(gw._scf.mo_coeff) nocca, noccb = gw.nocc nmoa, nmob = gw.nmo nkpts = gw.nkpts kpts = gw.kpts nklist = len(kptlist) nw = len(freqs) norbs = len(orbs) mydf = gw.with_df # possible kpts shift kscaled = gw.mol.get_scaled_kpts(kpts) kscaled -= kscaled[0] # This code does not support metals h**o = -99. lumo = 99. for k in range(nkpts): if h**o < max(mo_energy[0, k][nocca - 1], mo_energy[1, k][noccb - 1]): h**o = max(mo_energy[0, k][nocca - 1], mo_energy[1, k][noccb - 1]) if lumo > min(mo_energy[0, k][nocca], mo_energy[1, k][noccb]): lumo = min(mo_energy[0, k][nocca], mo_energy[1, k][noccb]) if (lumo - h**o) < 1e-3: logger.warn(gw, 'Current KUGW is not supporting metals!') ef = (h**o + lumo) / 2. # Integration on numerical grids if iw_cutoff is not None: nw_sigma = sum(iw < iw_cutoff for iw in freqs) + 1 else: nw_sigma = nw + 1 # Compute occ for -iw and vir for iw separately # to avoid branch cuts in analytic continuation omega_occ = np.zeros((nw_sigma), dtype=np.complex128) omega_vir = np.zeros((nw_sigma), dtype=np.complex128) omega_occ[0] = 1j * 0. omega_occ[1:] = -1j * freqs[:(nw_sigma - 1)] omega_vir[0] = 1j * 0. omega_vir[1:] = 1j * freqs[:(nw_sigma - 1)] orbs_occ_a = [i for i in orbs if i < nocca] orbs_occ_b = [i for i in orbs if i < noccb] norbs_occ_a = len(orbs_occ_a) norbs_occ_b = len(orbs_occ_b) emo_occ_a = np.zeros((nkpts, nmoa, nw_sigma), dtype=np.complex128) emo_occ_b = np.zeros((nkpts, nmob, nw_sigma), dtype=np.complex128) emo_vir_a = np.zeros((nkpts, nmoa, nw_sigma), dtype=np.complex128) emo_vir_b = np.zeros((nkpts, nmob, nw_sigma), dtype=np.complex128) for k in range(nkpts): emo_occ_a[k] = omega_occ[None, :] + ef - mo_energy[0, k][:, None] emo_occ_b[k] = omega_occ[None, :] + ef - mo_energy[1, k][:, None] emo_vir_a[k] = omega_vir[None, :] + ef - mo_energy[0, k][:, None] emo_vir_b[k] = omega_vir[None, :] + ef - mo_energy[1, k][:, None] sigma = np.zeros((2, nklist, norbs, nw_sigma), dtype=np.complex128) omega = np.zeros((2, norbs, nw_sigma), dtype=np.complex128) for s in range(2): for p in range(norbs): orbp = orbs[p] if orbp < gw.nocc[s]: omega[s, p] = omega_occ.copy() else: omega[s, p] = omega_vir.copy() if gw.fc: # Set up q mesh for q->0 finite size correction q_pts = np.array([1e-3, 0, 0]).reshape(1, 3) nq_pts = len(q_pts) q_abs = gw.mol.get_abs_kpts(q_pts) # Get qij = 1/sqrt(Omega) * < psi_{ik} | e^{iqr} | psi_{ak-q} > at q: (nkpts, nocc, nvir) qij = get_qij(gw, q_abs[0], mo_coeff) for kL in range(nkpts): # Lij: (2, ki, L, i, j) for looping every kL #Lij = np.zeros((2,nkpts,naux,nmoa,nmoa),dtype=np.complex128) Lij = [] # kidx: save kj that conserves with kL and ki (-ki+kj+kL=G) # kidx_r: save ki that conserves with kL and kj (-ki+kj+kL=G) kidx = np.zeros((nkpts), dtype=np.int64) kidx_r = np.zeros((nkpts), dtype=np.int64) for i, kpti in enumerate(kpts): for j, kptj in enumerate(kpts): # Find (ki,kj) that satisfies momentum conservation with kL kconserv = -kscaled[i] + kscaled[j] + kscaled[kL] is_kconserv = np.linalg.norm(np.round(kconserv) - kconserv) < 1e-12 if is_kconserv: kidx[i] = j kidx_r[j] = i logger.debug( gw, "Read Lpq (kL: %s / %s, ki: %s, kj: %s)" % (kL + 1, nkpts, i, j)) Lij_out_a = None Lij_out_b = None # Read (L|pq) and ao2mo transform to (L|ij) Lpq = [] for LpqR, LpqI, sign in mydf.sr_loop([kpti, kptj], max_memory=0.1 * gw._scf.max_memory, compact=False): Lpq.append(LpqR + LpqI * 1.0j) Lpq = np.vstack(Lpq).reshape(-1, nmoa**2) moija, ijslicea = _conc_mos(mo_coeff[0, i], mo_coeff[0, j])[2:] moijb, ijsliceb = _conc_mos(mo_coeff[1, i], mo_coeff[1, j])[2:] tao = [] ao_loc = None Lij_out_a = _ao2mo.r_e2(Lpq, moija, ijslicea, tao, ao_loc, out=Lij_out_a) tao = [] ao_loc = None Lij_out_b = _ao2mo.r_e2(Lpq, moijb, ijsliceb, tao, ao_loc, out=Lij_out_b) Lij.append( np.asarray((Lij_out_a.reshape(-1, nmoa, nmoa), Lij_out_b.reshape(-1, nmob, nmob)))) Lij = np.asarray(Lij) Lij = Lij.transpose(1, 0, 2, 3, 4) naux = Lij.shape[2] if kL == 0: for w in range(nw): # body dielectric matrix eps_body Pi = get_rho_response(gw, freqs[w], mo_energy, Lij, kL, kidx) eps_body_inv = np.linalg.inv(np.eye(naux) - Pi) if gw.fc: # head dielectric matrix eps_00 Pi_00 = get_rho_response_head(gw, freqs[w], mo_energy, qij) eps_00 = 1. - 4. * np.pi / np.linalg.norm( q_abs[0])**2 * Pi_00 # wings dielectric matrix eps_P0 Pi_P0 = get_rho_response_wing(gw, freqs[w], mo_energy, Lij, qij) eps_P0 = -np.sqrt(4. * np.pi) / np.linalg.norm( q_abs[0]) * Pi_P0 # inverse dielectric matrix eps_inv_00 = 1. / (eps_00 - np.dot( np.dot(eps_P0.conj(), eps_body_inv), eps_P0)) eps_inv_P0 = -eps_inv_00 * np.dot(eps_body_inv, eps_P0) # head correction Del_00 = 2. / np.pi * (6. * np.pi**2 / gw.mol.vol / nkpts )**(1. / 3.) * (eps_inv_00 - 1.) eps_inv_PQ = eps_body_inv g0_occ_a = wts[w] * emo_occ_a / (emo_occ_a**2 + freqs[w]**2) g0_occ_b = wts[w] * emo_occ_b / (emo_occ_b**2 + freqs[w]**2) g0_vir_a = wts[w] * emo_vir_a / (emo_vir_a**2 + freqs[w]**2) g0_vir_b = wts[w] * emo_vir_b / (emo_vir_b**2 + freqs[w]**2) for k in range(nklist): kn = kptlist[k] # Find km that conserves with kn and kL (-km+kn+kL=G) km = kidx_r[kn] Qmn_a = einsum('Pmn,PQ->Qmn', Lij[0, km][:, :, orbs].conj(), eps_inv_PQ - np.eye(naux)) Qmn_b = einsum('Pmn,PQ->Qmn', Lij[1, km][:, :, orbs].conj(), eps_inv_PQ - np.eye(naux)) Wmn_a = 1. / nkpts * einsum('Qmn,Qmn->mn', Qmn_a, Lij[0, km][:, :, orbs]) Wmn_b = 1. / nkpts * einsum('Qmn,Qmn->mn', Qmn_b, Lij[1, km][:, :, orbs]) sigma[0, k][:norbs_occ_a] += -einsum( 'mn,mw->nw', Wmn_a[:, :norbs_occ_a], g0_occ_a[km]) / np.pi sigma[1, k][:norbs_occ_b] += -einsum( 'mn,mw->nw', Wmn_b[:, :norbs_occ_b], g0_occ_b[km]) / np.pi sigma[0, k][norbs_occ_a:] += -einsum( 'mn,mw->nw', Wmn_a[:, norbs_occ_a:], g0_vir_a[km]) / np.pi sigma[1, k][norbs_occ_b:] += -einsum( 'mn,mw->nw', Wmn_b[:, norbs_occ_b:], g0_vir_b[km]) / np.pi if gw.fc: # apply head correction assert (kn == km) sigma[0, k][:norbs_occ_a] += -Del_00 * g0_occ_a[kn][ orbs][:norbs_occ_a] / np.pi sigma[0, k][norbs_occ_a:] += -Del_00 * g0_vir_a[kn][ orbs][norbs_occ_a:] / np.pi sigma[1, k][:norbs_occ_b] += -Del_00 * g0_occ_b[kn][ orbs][:norbs_occ_b] / np.pi sigma[1, k][norbs_occ_b:] += -Del_00 * g0_vir_b[kn][ orbs][norbs_occ_b:] / np.pi # apply wing correction Wn_P0_a = einsum('Pnm,P->nm', Lij[0, kn], eps_inv_P0).diagonal() Wn_P0_b = einsum('Pnm,P->nm', Lij[1, kn], eps_inv_P0).diagonal() Wn_P0_a = Wn_P0_a.real * 2. Wn_P0_b = Wn_P0_b.real * 2. Del_P0_a = np.sqrt(gw.mol.vol / 4. / np.pi**3) * ( 6. * np.pi**2 / gw.mol.vol / nkpts)**(2. / 3.) * Wn_P0_a[orbs] Del_P0_b = np.sqrt(gw.mol.vol / 4. / np.pi**3) * ( 6. * np.pi**2 / gw.mol.vol / nkpts)**(2. / 3.) * Wn_P0_b[orbs] sigma[0, k][:norbs_occ_a] += -einsum( 'n,nw->nw', Del_P0_a[:norbs_occ_a], g0_occ_a[kn][orbs][:norbs_occ_a]) / np.pi sigma[0, k][norbs_occ_a:] += -einsum( 'n,nw->nw', Del_P0_a[norbs_occ_a:], g0_vir_a[kn][orbs][norbs_occ_a:]) / np.pi sigma[1, k][:norbs_occ_b] += -einsum( 'n,nw->nw', Del_P0_b[:norbs_occ_b], g0_occ_b[kn][orbs][:norbs_occ_b]) / np.pi sigma[1, k][norbs_occ_b:] += -einsum( 'n,nw->nw', Del_P0_b[norbs_occ_b:], g0_vir_b[kn][orbs][norbs_occ_b:]) / np.pi else: for w in range(nw): Pi = get_rho_response(gw, freqs[w], mo_energy, Lij, kL, kidx) Pi_inv = np.linalg.inv(np.eye(naux) - Pi) - np.eye(naux) g0_occ_a = wts[w] * emo_occ_a / (emo_occ_a**2 + freqs[w]**2) g0_occ_b = wts[w] * emo_occ_b / (emo_occ_b**2 + freqs[w]**2) g0_vir_a = wts[w] * emo_vir_a / (emo_vir_a**2 + freqs[w]**2) g0_vir_b = wts[w] * emo_vir_b / (emo_vir_b**2 + freqs[w]**2) for k in range(nklist): kn = kptlist[k] # Find km that conserves with kn and kL (-km+kn+kL=G) km = kidx_r[kn] Qmn_a = einsum('Pmn,PQ->Qmn', Lij[0, km][:, :, orbs].conj(), Pi_inv) Qmn_b = einsum('Pmn,PQ->Qmn', Lij[1, km][:, :, orbs].conj(), Pi_inv) Wmn_a = 1. / nkpts * einsum('Qmn,Qmn->mn', Qmn_a, Lij[0, km][:, :, orbs]) Wmn_b = 1. / nkpts * einsum('Qmn,Qmn->mn', Qmn_b, Lij[1, km][:, :, orbs]) sigma[0, k][:norbs_occ_a] += -einsum( 'mn,mw->nw', Wmn_a[:, :norbs_occ_a], g0_occ_a[km]) / np.pi sigma[1, k][:norbs_occ_b] += -einsum( 'mn,mw->nw', Wmn_b[:, :norbs_occ_b], g0_occ_b[km]) / np.pi sigma[0, k][norbs_occ_a:] += -einsum( 'mn,mw->nw', Wmn_a[:, norbs_occ_a:], g0_vir_a[km]) / np.pi sigma[1, k][norbs_occ_b:] += -einsum( 'mn,mw->nw', Wmn_b[:, norbs_occ_b:], g0_vir_b[km]) / np.pi return sigma, omega
def general(mydf, mo_coeffs, kpts=None, compact=True): cell = mydf.cell kptijkl = _format_kpts(kpts) kpti, kptj, kptk, kptl = kptijkl if isinstance(mo_coeffs, numpy.ndarray) and mo_coeffs.ndim == 2: mo_coeffs = (mo_coeffs,) * 4 all_real = not any(numpy.iscomplexobj(mo) for mo in mo_coeffs) max_memory = max(2000, (mydf.max_memory - lib.current_memory()[0]) * .5) #################### # gamma point, the integral is real and with s4 symmetry if abs(kptijkl).sum() < KPT_DIFF_TOL and all_real: ijmosym, nij_pair, moij, ijslice = _conc_mos(mo_coeffs[0], mo_coeffs[1], compact) klmosym, nkl_pair, mokl, klslice = _conc_mos(mo_coeffs[2], mo_coeffs[3], compact) eri_mo = numpy.zeros((nij_pair,nkl_pair)) sym = (iden_coeffs(mo_coeffs[0], mo_coeffs[2]) and iden_coeffs(mo_coeffs[1], mo_coeffs[3])) coulG = mydf.weighted_coulG(kptj-kpti, False, mydf.gs) ijR = ijI = klR = klI = buf = None for pqkR, pqkI, p0, p1 \ in mydf.pw_loop(mydf.gs, kptijkl[:2], max_memory=max_memory, aosym='s2'): vG = numpy.sqrt(coulG[p0:p1]) pqkR *= vG pqkI *= vG buf = lib.transpose(pqkR, out=buf) ijR, klR = _dtrans(buf, ijR, ijmosym, moij, ijslice, buf, klR, klmosym, mokl, klslice, sym) lib.ddot(ijR.T, klR, 1, eri_mo, 1) buf = lib.transpose(pqkI, out=buf) ijI, klI = _dtrans(buf, ijI, ijmosym, moij, ijslice, buf, klI, klmosym, mokl, klslice, sym) lib.ddot(ijI.T, klI, 1, eri_mo, 1) pqkR = pqkI = None return eri_mo #################### # (kpt) i == j == k == l != 0 # (kpt) i == l && j == k && i != j && j != k => # elif (abs(kpti-kptl).sum() < KPT_DIFF_TOL) and (abs(kptj-kptk).sum() < KPT_DIFF_TOL): mo_coeffs = _mo_as_complex(mo_coeffs) nij_pair, moij, ijslice = _conc_mos(mo_coeffs[0], mo_coeffs[1])[1:] nlk_pair, molk, lkslice = _conc_mos(mo_coeffs[3], mo_coeffs[2])[1:] eri_mo = numpy.zeros((nij_pair,nlk_pair), dtype=numpy.complex) sym = (iden_coeffs(mo_coeffs[0], mo_coeffs[3]) and iden_coeffs(mo_coeffs[1], mo_coeffs[2])) coulG = mydf.weighted_coulG(kptj-kpti, False, mydf.gs) zij = zlk = buf = None for pqkR, pqkI, p0, p1 \ in mydf.pw_loop(mydf.gs, kptijkl[:2], max_memory=max_memory): buf = lib.transpose(pqkR+pqkI*1j, out=buf) buf *= numpy.sqrt(coulG[p0:p1]).reshape(-1,1) zij, zlk = _ztrans(buf, zij, moij, ijslice, buf, zlk, molk, lkslice, sym) lib.dot(zij.T, zlk.conj(), 1, eri_mo, 1) pqkR = pqkI = None nmok = mo_coeffs[2].shape[1] nmol = mo_coeffs[3].shape[1] eri_mo = lib.transpose(eri_mo.reshape(-1,nmol,nmok), axes=(0,2,1)) return eri_mo.reshape(nij_pair,nlk_pair) #################### # aosym = s1, complex integrals # # If kpti == kptj, (kptl-kptk)*a has to be multiples of 2pi because of the wave # vector symmetry. k is a fraction of reciprocal basis, 0 < k/b < 1, by definition. # So kptl/b - kptk/b must be -1 < k/b < 1. => kptl == kptk # else: mo_coeffs = _mo_as_complex(mo_coeffs) nij_pair, moij, ijslice = _conc_mos(mo_coeffs[0], mo_coeffs[1])[1:] nkl_pair, mokl, klslice = _conc_mos(mo_coeffs[2], mo_coeffs[3])[1:] eri_mo = numpy.zeros((nij_pair,nkl_pair), dtype=numpy.complex) tao = [] ao_loc = None coulG = mydf.weighted_coulG(kptj-kpti, False, mydf.gs) zij = zkl = buf = None for (pqkR, pqkI, p0, p1), (rskR, rskI, q0, q1) in \ lib.izip(mydf.pw_loop(mydf.gs, kptijkl[:2], max_memory=max_memory*.5), mydf.pw_loop(mydf.gs,-kptijkl[2:], max_memory=max_memory*.5)): buf = lib.transpose(pqkR+pqkI*1j, out=buf) zij = _ao2mo.r_e2(buf, moij, ijslice, tao, ao_loc, out=zij) buf = lib.transpose(rskR-rskI*1j, out=buf) zkl = _ao2mo.r_e2(buf, mokl, klslice, tao, ao_loc, out=zkl) zij *= coulG[p0:p1].reshape(-1,1) lib.dot(zij.T, zkl, 1, eri_mo, 1) pqkR = pqkI = rskR = rskI = None return eri_mo