def get_mo_pairs_G(mydf, mo_coeffs, kpts=numpy.zeros((2,3)), q=None, compact=False): '''Calculate forward (G|ij) FFT of all MO pairs. Args: mo_coeff: length-2 list of (nao,nmo) ndarrays The two sets of MO coefficients to use in calculating the product |ij). Returns: mo_pairs_G : (ngs, nmoi*nmoj) ndarray The FFT of the real-space MO pairs. ''' if kpts is None: kpts = numpy.zeros((2,3)) cell = mydf.cell kpts = numpy.asarray(kpts) coords = cell.gen_uniform_grids(mydf.gs) nmoi = mo_coeffs[0].shape[1] nmoj = mo_coeffs[1].shape[1] ngs = len(coords) def trans(aoiR, aojR, fac=1): if id(aoiR) == id(aojR) and iden_coeffs(mo_coeffs[0], mo_coeffs[1]): moiR = mojR = numpy.asarray(lib.dot(mo_coeffs[0].T,aoiR.T), order='C') else: moiR = numpy.asarray(lib.dot(mo_coeffs[0].T, aoiR.T), order='C') mojR = numpy.asarray(lib.dot(mo_coeffs[1].T, aojR.T), order='C') mo_pairs_G = numpy.empty((nmoi,nmoj,ngs), dtype=numpy.complex128) for i in range(nmoi): mo_pairs_G[i] = tools.fft(fac * moiR[i].conj() * mojR, mydf.gs) mo_pairs_G = mo_pairs_G.reshape(-1,ngs).T return mo_pairs_G if gamma_point(kpts): # gamma point, real aoR = mydf._numint.eval_ao(cell, coords, kpts[:1])[0] if compact and iden_coeffs(mo_coeffs[0], mo_coeffs[1]): moR = numpy.asarray(lib.dot(mo_coeffs[0].T, aoR.T), order='C') npair = nmoi*(nmoi+1)//2 mo_pairs_G = numpy.empty((npair,ngs), dtype=numpy.complex128) ij = 0 for i in range(nmoi): mo_pairs_G[ij:ij+i+1] = tools.fft(moR[i].conj() * moR[:i+1], mydf.gs) ij += i + 1 mo_pairs_G = mo_pairs_G.T else: mo_pairs_G = trans(aoR, aoR) elif is_zero(kpts[0]-kpts[1]): aoR = mydf._numint.eval_ao(cell, coords, kpts[:1])[0] mo_pairs_G = trans(aoR, aoR) else: if q is None: q = kpts[1] - kpts[0] aoiR, aojR = mydf._numint.eval_ao(cell, coords, kpts) fac = numpy.exp(-1j * numpy.dot(coords, q)) mo_pairs_G = trans(aoiR, aojR, fac) return mo_pairs_G
def get_j_for_bands(mydf, dm_kpts, hermi=1, kpts=numpy.zeros((1, 3)), kpts_band=None): cell = mydf.cell log = logger.Logger(mydf.stdout, mydf.verbose) t1 = (time.clock(), time.time()) dm_kpts = lib.asarray(dm_kpts, order='C') dms = _format_dms(dm_kpts, kpts) nset, nkpts, nao = dms.shape[:3] dmsR = dms.real.reshape(nset, nkpts, nao**2) dmsI = dms.imag.reshape(nset, nkpts, nao**2) kpt_allow = numpy.zeros(3) coulG = mydf.weighted_coulG(kpt_allow, False, mydf.gs) ngs = len(coulG) vG = numpy.zeros((nset, ngs), dtype=numpy.complex128) max_memory = (mydf.max_memory - lib.current_memory()[0]) * .8 dmsC = dms.conj() for aoaoks, p0, p1 in mydf.ft_loop(mydf.gs, kpt_allow, kpts, max_memory=max_memory): #:rho = numpy.einsum('lkL,lk->L', pqk.conj(), dm) for k, aoao in enumerate(aoaoks): for i in range(nset): rho = numpy.einsum('ij,Lij->L', dmsC[i, k], aoao).conj() vG[i, p0:p1] += rho * coulG[p0:p1] aoao = aoaoks = p0 = p1 = None weight = 1. / len(kpts) vG *= weight t1 = log.timer_debug1('get_j pass 1 to compute J(G)', *t1) kpts_band, single_kpt_band = _format_kpts_band(kpts_band, kpts) nband = len(kpts_band) vj_kpts = numpy.zeros((nset, nband, nao, nao), dtype=numpy.complex128) for aoaoks, p0, p1 in mydf.ft_loop(mydf.gs, kpt_allow, kpts_band, max_memory=max_memory): for k, aoao in enumerate(aoaoks): for i in range(nset): vj_kpts[i, k] += numpy.einsum('L,Lij->ij', vG[i, p0:p1], aoao) aoao = aoaoks = p0 = p1 = None if gamma_point(kpts_band): vj_kpts = vj_kpts.real.copy() t1 = log.timer_debug1('get_j pass 2', *t1) return _format_jks(vj_kpts, dm_kpts, kpts_band, kpts, single_kpt_band)
def ft_ao(mol, Gv, shls_slice=None, b=None, gxyz=None, Gvbase=None, kpt=numpy.zeros(3), verbose=None): if gamma_point(kpt): return mol_ft_ao(mol, Gv, shls_slice, b, gxyz, Gvbase, verbose) else: kG = Gv + kpt return mol_ft_ao(mol, kG, shls_slice, None, None, None, verbose)
def general(mydf, mo_coeffs, kpts=None, compact=False): '''General MO integral transformation''' 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 mo_coeffs = [numpy.asarray(mo, order='F') for mo in mo_coeffs] allreal = not any(numpy.iscomplexobj(mo) for mo in mo_coeffs) q = kptj - kpti coulG = tools.get_coulG(cell, q, gs=mydf.gs) coords = cell.gen_uniform_grids(mydf.gs) max_memory = mydf.max_memory - lib.current_memory()[0] if gamma_point(kptijkl) and allreal: ao = mydf._numint.eval_ao(cell, coords, kpti)[0] if ((iden_coeffs(mo_coeffs[0], mo_coeffs[2]) and iden_coeffs(mo_coeffs[1], mo_coeffs[3]))): moiT = mojT = numpy.asarray(lib.dot(mo_coeffs[0].T, ao.T), order='C') ao = None max_memory = max_memory - moiT.nbytes * 1e-6 eri = _contract_compact(mydf, (moiT, mojT), coulG, max_memory=max_memory) if not compact: nmo = moiT.shape[0] eri = ao2mo.restore(1, eri, nmo).reshape(nmo**2, nmo**2) else: mos = [ numpy.asarray(lib.dot(c.T, ao.T), order='C') for c in mo_coeffs ] ao = None fac = numpy.array(1.) max_memory = max_memory - sum([x.nbytes for x in mos]) * 1e-6 eri = _contract_plain(mydf, mos, coulG, fac, max_memory=max_memory).real return eri else: aos = mydf._numint.eval_ao(cell, coords, kptijkl) mos = [ numpy.asarray(lib.dot(c.T, aos[i].T), order='C') for i, c in enumerate(mo_coeffs) ] aos = None fac = numpy.exp(-1j * numpy.dot(coords, q)) max_memory = max_memory - sum([x.nbytes for x in mos]) * 1e-6 eri = _contract_plain(mydf, mos, coulG, fac, max_memory=max_memory) return eri
def get_j_kpts(mydf, dm_kpts, hermi=1, kpts=np.zeros((1, 3)), kpts_band=None): '''Get the Coulomb (J) AO matrix at sampled k-points. Args: dm_kpts : (nkpts, nao, nao) ndarray or a list of (nkpts,nao,nao) ndarray Density matrix at each k-point. If a list of k-point DMs, eg, UHF alpha and beta DM, the alpha and beta DMs are contracted separately. kpts : (nkpts, 3) ndarray Kwargs: kpts_band : (3,) ndarray or (*,3) ndarray A list of arbitrary "band" k-points at which to evalute the matrix. Returns: vj : (nkpts, nao, nao) ndarray or list of vj if the input dm_kpts is a list of DMs ''' cell = mydf.cell gs = mydf.gs dm_kpts = lib.asarray(dm_kpts, order='C') dms = _format_dms(dm_kpts, kpts) nset, nkpts, nao = dms.shape[:3] coulG = tools.get_coulG(cell, gs=gs) ngs = len(coulG) vR = rhoR = np.zeros((nset, ngs)) for k, aoR in mydf.aoR_loop(gs, kpts): for i in range(nset): rhoR[i] += numint.eval_rho(cell, aoR, dms[i, k]) for i in range(nset): rhoR[i] *= 1. / nkpts rhoG = tools.fft(rhoR[i], gs) vG = coulG * rhoG vR[i] = tools.ifft(vG, gs).real kpts_band, input_band = _format_kpts_band(kpts_band, kpts), kpts_band nband = len(kpts_band) weight = cell.vol / ngs if gamma_point(kpts_band): vj_kpts = np.empty((nset, nband, nao, nao)) else: vj_kpts = np.empty((nset, nband, nao, nao), dtype=np.complex128) for k, aoR in mydf.aoR_loop(gs, kpts_band): for i in range(nset): vj_kpts[i, k] = weight * lib.dot(aoR.T.conj() * vR[i], aoR) return _format_jks(vj_kpts, dm_kpts, input_band, kpts)
def get_eri(mydf, kpts=None, compact=False): cell = mydf.cell kptijkl = _format_kpts(kpts) kpti, kptj, kptk, kptl = kptijkl q = kptj - kpti coulG = tools.get_coulG(cell, q, gs=mydf.gs) coords = cell.gen_uniform_grids(mydf.gs) max_memory = mydf.max_memory - lib.current_memory()[0] #################### # gamma point, the integral is real and with s4 symmetry if gamma_point(kptijkl): #:ao_pairs_G = get_ao_pairs_G(mydf, kptijkl[:2], q, compact=compact) #:ao_pairs_G *= numpy.sqrt(coulG).reshape(-1,1) #:eri = lib.dot(ao_pairs_G.T, ao_pairs_G, cell.vol/ngs**2) ao = mydf._numint.eval_ao(cell, coords, kpti)[0] ao = numpy.asarray(ao.T, order='C') eri = _contract_compact(mydf, (ao, ao), coulG, max_memory=max_memory) if not compact: nao = cell.nao_nr() eri = ao2mo.restore(1, eri, nao).reshape(nao**2, nao**2) return eri #################### # aosym = s1, complex integrals else: #:ao_pairs_G = get_ao_pairs_G(mydf, kptijkl[:2], q, compact=False) #:# ao_pairs_invG = rho_kl(-(G+k_ij)) = conj(rho_lk(G+k_ij)).swap(r,s) #:#=get_ao_pairs_G(mydf, [kptl,kptk], q, compact=False).transpose(0,2,1).conj() #:ao_pairs_invG = get_ao_pairs_G(mydf, -kptijkl[2:], q, compact=False).conj() #:ao_pairs_G *= coulG.reshape(-1,1) #:eri = lib.dot(ao_pairs_G.T, ao_pairs_invG, cell.vol/ngs**2) if is_zero(kpti - kptl) and is_zero(kptj - kptk): if is_zero(kpti - kptj): aoi = mydf._numint.eval_ao(cell, coords, kpti)[0] aoi = aoj = numpy.asarray(aoi.T, order='C') else: aoi, aoj = mydf._numint.eval_ao(cell, coords, kptijkl[:2]) aoi = numpy.asarray(aoi.T, order='C') aoj = numpy.asarray(aoj.T, order='C') aos = (aoi, aoj, aoj, aoi) else: aos = mydf._numint.eval_ao(cell, coords, kptijkl) aos = [numpy.asarray(x.T, order='C') for x in aos] fac = numpy.exp(-1j * numpy.dot(coords, q)) max_memory = max_memory - aos[0].nbytes * 4 * 1e-6 eri = _contract_plain(mydf, aos, coulG, fac, max_memory=max_memory) return eri
def aux_e2(cell, auxcell, intor='int3c2e_sph', aosym='s1', comp=1, kptij_lst=numpy.zeros((1, 2, 3)), shls_slice=None): '''3-center AO integrals (ij|L) with double lattice sum: \sum_{lm} (i[l]j[m]|L[0]), where L is the auxiliary basis. Returns: (nao_pair, naux) array ''' intor = gto.moleintor.ascint3(intor) if shls_slice is None: shls_slice = (0, cell.nbas, 0, cell.nbas, 0, auxcell.nbas) ao_loc = cell.ao_loc_nr() aux_loc = auxcell.ao_loc_nr('ssc' in intor)[:shls_slice[5] + 1] ni = ao_loc[shls_slice[1]] - ao_loc[shls_slice[0]] nj = ao_loc[shls_slice[3]] - ao_loc[shls_slice[2]] naux = aux_loc[shls_slice[5]] - aux_loc[shls_slice[4]] nkptij = len(kptij_lst) kpti = kptij_lst[:, 0] kptj = kptij_lst[:, 1] j_only = is_zero(kpti - kptj) if j_only and aosym[:2] == 's2': assert (shls_slice[2] == 0) nao_pair = (ao_loc[shls_slice[1]] * (ao_loc[shls_slice[1]] + 1) // 2 - ao_loc[shls_slice[0]] * (ao_loc[shls_slice[0]] + 1) // 2) else: nao_pair = ni * nj if gamma_point(kptij_lst): dtype = numpy.double else: dtype = numpy.complex128 int3c = wrap_int3c(cell, auxcell, intor, aosym, comp, kptij_lst) out = numpy.empty((nkptij, comp, nao_pair, naux), dtype=dtype) out = int3c(shls_slice, out) if comp == 1: out = out[:, 0] if nkptij == 1: out = out[0] return out
def get_eri(mydf, kpts=None, compact=False): cell = mydf.cell kptijkl = _format_kpts(kpts) kpti, kptj, kptk, kptl = kptijkl nao = cell.nao_nr() nao_pair = nao * (nao+1) // 2 q = kptj - kpti coulG = tools.get_coulG(cell, q, gs=mydf.gs) ngs = len(coulG) #################### # gamma point, the integral is real and with s4 symmetry if gamma_point(kptijkl): ao_pairs_G = get_ao_pairs_G(mydf, kptijkl[:2], q, compact=compact) ao_pairs_G *= numpy.sqrt(coulG).reshape(-1,1) aoijR = ao_pairs_G.real.copy() aoijI = ao_pairs_G.imag.copy() ao_pairs_G = None eri = lib.dot(aoijR.T, aoijR, cell.vol/ngs**2) eri = lib.dot(aoijI.T, aoijI, cell.vol/ngs**2, eri, 1) return eri #################### # (kpt) i == j == k == l != 0 # (kpt) i == l && j == k && i != j && j != k => # # complex integrals, N^4 elements elif is_zero(kpti-kptl) and is_zero(kptj-kptk): ao_pairs_G = get_ao_pairs_G(mydf, kptijkl[:2], q, compact=False) ao_pairs_G *= numpy.sqrt(coulG).reshape(-1,1) ao_pairs_invG = ao_pairs_G.T.reshape(nao,nao,-1).transpose(1,0,2).conj() ao_pairs_invG = ao_pairs_invG.reshape(-1,ngs) return lib.dot(ao_pairs_G.T, ao_pairs_invG.T, cell.vol/ngs**2) #################### # aosym = s1, complex integrals # else: ao_pairs_G = get_ao_pairs_G(mydf, kptijkl[:2], q, compact=False) # ao_pairs_invG = rho_rs(-G+k_rs) = conj(rho_sr(G+k_sr)).swap(r,s) ao_pairs_invG = get_ao_pairs_G(mydf, -kptijkl[2:], q, compact=False).conj() ao_pairs_G *= coulG.reshape(-1,1) return lib.dot(ao_pairs_G.T, ao_pairs_invG, cell.vol/ngs**2)
def get_j_kpts(mydf, dm_kpts, hermi=1, kpts=numpy.zeros((1, 3)), kpts_band=None): if kpts_band is not None: return get_j_for_bands(mydf, dm_kpts, hermi, kpts, kpts_band) cell = mydf.cell log = logger.Logger(mydf.stdout, mydf.verbose) t1 = (time.clock(), time.time()) dm_kpts = lib.asarray(dm_kpts, order='C') dms = _format_dms(dm_kpts, kpts) nset, nkpts, nao = dms.shape[:3] vj_kpts = numpy.zeros((nset, nkpts, nao, nao), dtype=numpy.complex128) kpt_allow = numpy.zeros(3) coulG = mydf.weighted_coulG(kpt_allow, False, mydf.gs) max_memory = (mydf.max_memory - lib.current_memory()[0]) * .8 weight = 1. / len(kpts) dmsC = dms.conj() for aoaoks, p0, p1 in mydf.ft_loop(mydf.gs, kpt_allow, kpts, max_memory=max_memory): vG = [0] * nset #:rho = numpy.einsum('lkL,lk->L', pqk.conj(), dm) for k, aoao in enumerate(aoaoks): for i in range(nset): rho = numpy.einsum('ij,Lij->L', dmsC[i, k], aoao).conj() vG[i] += rho * coulG[p0:p1] for i in range(nset): vG[i] *= weight for k, aoao in enumerate(aoaoks): for i in range(nset): vj_kpts[i, k] += numpy.einsum('L,Lij->ij', vG[i], aoao) aoao = aoaoks = p0 = p1 = None if gamma_point(kpts): vj_kpts = vj_kpts.real.copy() return _format_jks(vj_kpts, dm_kpts, kpts_band, kpts)
def get_nuc(mydf, kpts=None): cell = mydf.cell if kpts is None: kpts_lst = numpy.zeros((1, 3)) else: kpts_lst = numpy.reshape(kpts, (-1, 3)) log = logger.Logger(mydf.stdout, mydf.verbose) t1 = (time.clock(), time.time()) nkpts = len(kpts_lst) nao = cell.nao_nr() nao_pair = nao * (nao + 1) // 2 Gv, Gvbase, kws = cell.get_Gv_weights(mydf.gs) kpt_allow = numpy.zeros(3) if mydf.eta == 0: vpplocG = pseudo.pp_int.get_gth_vlocG_part1(cell, Gv) vpplocG = -numpy.einsum('ij,ij->j', cell.get_SI(Gv), vpplocG) vpplocG *= kws vG = vpplocG vj = numpy.zeros((nkpts, nao_pair), dtype=numpy.complex128) else: nuccell = copy.copy(cell) half_sph_norm = .5 / numpy.sqrt(numpy.pi) norm = half_sph_norm / gto.mole._gaussian_int(2, mydf.eta) chg_env = [mydf.eta, norm] ptr_eta = cell._env.size ptr_norm = ptr_eta + 1 chg_bas = [[ia, 0, 1, 1, 0, ptr_eta, ptr_norm, 0] for ia in range(cell.natm)] nuccell._atm = cell._atm nuccell._bas = numpy.asarray(chg_bas, dtype=numpy.int32) nuccell._env = numpy.hstack((cell._env, chg_env)) # PP-loc part1 is handled by fakenuc in _int_nuc_vloc vj = lib.asarray(mydf._int_nuc_vloc(nuccell, kpts_lst)) t1 = log.timer_debug1('vnuc pass1: analytic int', *t1) charge = -cell.atom_charges() coulG = tools.get_coulG(cell, kpt_allow, gs=mydf.gs, Gv=Gv) coulG *= kws aoaux = ft_ao.ft_ao(nuccell, Gv) vG = numpy.einsum('i,xi->x', charge, aoaux) * coulG max_memory = max(2000, mydf.max_memory - lib.current_memory()[0]) for aoaoks, p0, p1 in mydf.ft_loop(mydf.gs, kpt_allow, kpts_lst, max_memory=max_memory, aosym='s2'): for k, aoao in enumerate(aoaoks): # rho_ij(G) nuc(-G) / G^2 # = [Re(rho_ij(G)) + Im(rho_ij(G))*1j] [Re(nuc(G)) - Im(nuc(G))*1j] / G^2 if gamma_point(kpts_lst[k]): vj[k] += numpy.einsum('k,kx->x', vG[p0:p1].real, aoao.real) vj[k] += numpy.einsum('k,kx->x', vG[p0:p1].imag, aoao.imag) else: vj[k] += numpy.einsum('k,kx->x', vG[p0:p1].conj(), aoao) t1 = log.timer_debug1('contracting Vnuc', *t1) vj_kpts = [] for k, kpt in enumerate(kpts_lst): if gamma_point(kpt): vj_kpts.append(lib.unpack_tril(vj[k].real.copy())) else: vj_kpts.append(lib.unpack_tril(vj[k])) if kpts is None or numpy.shape(kpts) == (3, ): vj_kpts = vj_kpts[0] return numpy.asarray(vj_kpts)
def get_jk(mydf, dm, hermi=1, kpt=numpy.zeros(3), kpts_band=None, with_j=True, with_k=True, exxdiv=None): '''JK for given k-point''' vj = vk = None if kpts_band is not None and abs(kpt - kpts_band).sum() > 1e-9: kpt = numpy.reshape(kpt, (1, 3)) if with_k: vk = get_k_kpts(mydf, dm, hermi, kpt, kpts_band, exxdiv) if with_j: vj = get_j_kpts(mydf, dm, hermi, kpt, kpts_band) return vj, vk cell = mydf.cell log = logger.Logger(mydf.stdout, mydf.verbose) t1 = (time.clock(), time.time()) if mydf._cderi is None or not mydf.has_kpts(kpts_band): if mydf._cderi is not None: log.warn( 'DF integrals for band k-points were not found %s. ' 'DF integrals will be rebuilt to include band k-points.', mydf._cderi) mydf.build(kpts_band=kpts_band) t1 = log.timer_debug1('Init get_jk', *t1) dm = numpy.asarray(dm, order='C') dms = _format_dms(dm, [kpt]) nset, _, nao = dms.shape[:3] dms = dms.reshape(nset, nao, nao) j_real = gamma_point(kpt) k_real = gamma_point(kpt) and not numpy.iscomplexobj(dms) kptii = numpy.asarray((kpt, kpt)) dmsR = dms.real.reshape(nset, nao, nao) dmsI = dms.imag.reshape(nset, nao, nao) mem_now = lib.current_memory()[0] max_memory = max(2000, (mydf.max_memory - mem_now)) if with_j: vjR = numpy.zeros((nset, nao, nao)) vjI = numpy.zeros((nset, nao, nao)) if with_k: vkR = numpy.zeros((nset, nao, nao)) vkI = numpy.zeros((nset, nao, nao)) buf1R = numpy.empty((mydf.blockdim * nao**2)) buf2R = numpy.empty((mydf.blockdim * nao**2)) buf1I = numpy.zeros((mydf.blockdim * nao**2)) buf2I = numpy.empty((mydf.blockdim * nao**2)) max_memory *= .5 log.debug1('max_memory = %d MB (%d in use)', max_memory, mem_now) def contract_k(pLqR, pLqI): # K ~ 'iLj,lLk*,li->kj' + 'lLk*,iLj,li->kj' #:pLq = (LpqR + LpqI.reshape(-1,nao,nao)*1j).transpose(1,0,2) #:tmp = numpy.dot(dm, pLq.reshape(nao,-1)) #:vk += numpy.dot(pLq.reshape(-1,nao).conj().T, tmp.reshape(-1,nao)) nrow = pLqR.shape[1] tmpR = numpy.ndarray((nao, nrow * nao), buffer=buf2R) if k_real: for i in range(nset): lib.ddot(dmsR[i], pLqR.reshape(nao, -1), 1, tmpR) lib.ddot( pLqR.reshape(-1, nao).T, tmpR.reshape(-1, nao), 1, vkR[i], 1) else: tmpI = numpy.ndarray((nao, nrow * nao), buffer=buf2I) for i in range(nset): zdotNN(dmsR[i], dmsI[i], pLqR.reshape(nao, -1), pLqI.reshape(nao, -1), 1, tmpR, tmpI, 0) zdotCN( pLqR.reshape(-1, nao).T, pLqI.reshape(-1, nao).T, tmpR.reshape(-1, nao), tmpI.reshape(-1, nao), 1, vkR[i], vkI[i], 1) pLqI = None thread_k = None for LpqR, LpqI in mydf.sr_loop(kptii, max_memory, False): LpqR = LpqR.reshape(-1, nao, nao) t1 = log.timer_debug1(' load', *t1) if thread_k is not None: thread_k.join() if with_j: #:rho_coeff = numpy.einsum('Lpq,xqp->xL', Lpq, dms) #:vj += numpy.dot(rho_coeff, Lpq.reshape(-1,nao**2)) rhoR = numpy.einsum('Lpq,xpq->xL', LpqR, dmsR) if not j_real: LpqI = LpqI.reshape(-1, nao, nao) rhoR -= numpy.einsum('Lpq,xpq->xL', LpqI, dmsI) rhoI = numpy.einsum('Lpq,xpq->xL', LpqR, dmsI) rhoI += numpy.einsum('Lpq,xpq->xL', LpqI, dmsR) vjR += numpy.einsum('xL,Lpq->xpq', rhoR, LpqR) if not j_real: vjR -= numpy.einsum('xL,Lpq->xpq', rhoI, LpqI) vjI += numpy.einsum('xL,Lpq->xpq', rhoR, LpqI) vjI += numpy.einsum('xL,Lpq->xpq', rhoI, LpqR) t1 = log.timer_debug1(' with_j', *t1) if with_k: nrow = LpqR.shape[0] pLqR = numpy.ndarray((nao, nrow, nao), buffer=buf1R) pLqR[:] = LpqR.transpose(1, 0, 2) if not k_real: pLqI = numpy.ndarray((nao, nrow, nao), buffer=buf1I) if LpqI is not None: pLqI[:] = LpqI.reshape(-1, nao, nao).transpose(1, 0, 2) thread_k = lib.background_thread(contract_k, pLqR, pLqI) t1 = log.timer_debug1(' with_k', *t1) LpqR = LpqI = pLqR = pLqI = None if thread_k is not None: thread_k.join() thread_k = None if with_j: if j_real: vj = vjR else: vj = vjR + vjI * 1j vj = vj.reshape(dm.shape) if with_k: if k_real: vk = vkR else: vk = vkR + vkI * 1j if exxdiv: assert (exxdiv.lower() == 'ewald') _ewald_exxdiv_for_G0(cell, kpt, dms, vk) vk = vk.reshape(dm.shape) t1 = log.timer('sr jk', *t1) return vj, vk
def get_eri(mydf, kpts=None, compact=True): if mydf._cderi is None or mydf.auxcell is None: mydf.build() cell = mydf.cell kptijkl = _format_kpts(kpts) kpti, kptj, kptk, kptl = kptijkl nao = cell.nao_nr() nao_pair = nao * (nao + 1) // 2 max_memory = max( 2000, mydf.max_memory - lib.current_memory()[0] - nao**4 * 8 / 1e6) #################### # gamma point, the integral is real and with s4 symmetry if gamma_point(kptijkl): eriR = numpy.zeros((nao_pair, nao_pair)) for LpqR, LpqI in mydf.sr_loop(kptijkl[:2], max_memory, True): lib.ddot(LpqR.T, LpqR, 1, eriR, 1) LpqR = LpqI = None if not compact: eriR = ao2mo.restore(1, eriR, nao).reshape(nao**2, -1) return eriR elif is_zero(kpti - kptk) and is_zero(kptj - kptl): eriR = numpy.zeros((nao * nao, nao * nao)) eriI = numpy.zeros((nao * nao, nao * nao)) for LpqR, LpqI in mydf.sr_loop(kptijkl[:2], max_memory, False): zdotNN(LpqR.T, LpqI.T, LpqR, LpqI, 1, eriR, eriI, 1) LpqR = LpqI = None return eriR + eriI * 1j #################### # (kpt) i == j == k == l != 0 # # (kpt) i == l && j == k && i != j && j != k => # both vbar and ovlp are zero. It corresponds to the exchange integral. # # complex integrals, N^4 elements elif is_zero(kpti - kptl) and is_zero(kptj - kptk): eriR = numpy.zeros((nao * nao, nao * nao)) eriI = numpy.zeros((nao * nao, nao * nao)) for LpqR, LpqI in mydf.sr_loop(kptijkl[:2], max_memory, False): zdotNC(LpqR.T, LpqI.T, LpqR, LpqI, 1, eriR, eriI, 1) LpqR = LpqI = None # transpose(0,1,3,2) because # j == k && i == l => # (L|ij).transpose(0,2,1).conj() = (L^*|ji) = (L^*|kl) => (M|kl) eri = lib.transpose((eriR + eriI * 1j).reshape(-1, nao, nao), axes=(0, 2, 1)) return eri.reshape(nao**2, -1) #################### # aosym = s1, complex integrals # # kpti == kptj => kptl == kptk # 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. # else: eriR = numpy.zeros((nao * nao, nao * nao)) eriI = numpy.zeros((nao * nao, nao * nao)) for (LpqR, LpqI), (LrsR, LrsI) in \ lib.izip(mydf.sr_loop(kptijkl[:2], max_memory, False), mydf.sr_loop(kptijkl[2:], max_memory, False)): zdotNN(LpqR.T, LpqI.T, LrsR, LrsI, 1, eriR, eriI, 1) LpqR = LpqI = LrsR = LrsI = None return eriR + eriI * 1j
def general(mydf, mo_coeffs, kpts=None, compact=True): if mydf._cderi is None or mydf.auxcell is None: mydf.build() 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 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 = klR = None for LpqR, LpqI in mydf.sr_loop(kptijkl[:2], max_memory, True): ijR, klR = _dtrans(LpqR, ijR, ijmosym, moij, ijslice, LpqR, klR, klmosym, mokl, klslice, sym) lib.ddot(ijR.T, klR, 1, eri_mo, 1) LpqR = LpqI = None return eri_mo elif is_zero(kpti - kptk) and is_zero(kptj - kptl): 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) sym = (iden_coeffs(mo_coeffs[0], mo_coeffs[2]) and iden_coeffs(mo_coeffs[1], mo_coeffs[3])) zij = zkl = None for LpqR, LpqI in mydf.sr_loop(kptijkl[:2], max_memory, False): buf = LpqR + LpqI * 1j zij, zkl = _ztrans(buf, zij, moij, ijslice, buf, zkl, mokl, klslice, sym) lib.dot(zij.T, zkl, 1, eri_mo, 1) LpqR = LpqI = buf = 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 = None for LpqR, LpqI in mydf.sr_loop(kptijkl[:2], max_memory, False): buf = LpqR + LpqI * 1j zij, zlk = _ztrans(buf, zij, moij, ijslice, buf, zlk, molk, lkslice, sym) lib.dot(zij.T, zlk.conj(), 1, eri_mo, 1) LpqR = LpqI = buf = 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) zij = zkl = None for (LpqR, LpqI), (LrsR, LrsI) in \ lib.izip(mydf.sr_loop(kptijkl[:2], max_memory, False), mydf.sr_loop(kptijkl[2:], max_memory, False)): zij, zkl = _ztrans(LpqR + LpqI * 1j, zij, moij, ijslice, LrsR + LrsI * 1j, zkl, mokl, klslice, False) lib.dot(zij.T, zkl, 1, eri_mo, 1) LpqR = LpqI = LrsR = LrsI = None return eri_mo
def get_eri(mydf, kpts=None, compact=True): cell = mydf.cell kptijkl = _format_kpts(kpts) kpti, kptj, kptk, kptl = kptijkl q = kptj - kpti coulG = mydf.weighted_coulG(q, False, mydf.gs) nao = cell.nao_nr() nao_pair = nao * (nao + 1) // 2 max_memory = max(2000, (mydf.max_memory - lib.current_memory()[0]) * .8) #################### # gamma point, the integral is real and with s4 symmetry if gamma_point(kptijkl): eriR = numpy.zeros((nao_pair, nao_pair)) 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 lib.ddot(pqkR, pqkR.T, 1, eriR, 1) lib.ddot(pqkI, pqkI.T, 1, eriR, 1) pqkR = pqkI = None if not compact: eriR = ao2mo.restore(1, eriR, nao).reshape(nao**2, -1) return eriR #################### # (kpt) i == j == k == l != 0 # (kpt) i == l && j == k && i != j && j != k => # # complex integrals, N^4 elements elif is_zero(kpti - kptl) and is_zero(kptj - kptk): eriR = numpy.zeros((nao**2, nao**2)) eriI = numpy.zeros((nao**2, nao**2)) for pqkR, pqkI, p0, p1 \ in mydf.pw_loop(mydf.gs, kptijkl[:2], q, max_memory=max_memory): vG = numpy.sqrt(coulG[p0:p1]) pqkR *= vG pqkI *= vG # rho_pq(G+k_pq) * conj(rho_rs(G-k_rs)) zdotNC(pqkR, pqkI, pqkR.T, pqkI.T, 1, eriR, eriI, 1) pqkR = pqkI = None pqkR = pqkI = coulG = None # transpose(0,1,3,2) because # j == k && i == l => # (L|ij).transpose(0,2,1).conj() = (L^*|ji) = (L^*|kl) => (M|kl) # rho_rs(-G+k_rs) = conj(transpose(rho_sr(G+k_sr), (0,2,1))) eri = lib.transpose((eriR + eriI * 1j).reshape(-1, nao, nao), axes=(0, 2, 1)) return eri.reshape(nao**2, -1) #################### # 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: eriR = numpy.zeros((nao**2, nao**2)) eriI = numpy.zeros((nao**2, nao**2)) # rho_rs(-G-k) = rho_rs(conj(G+k)) = conj(rho_sr(G+k)) 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)): pqkR *= coulG[p0:p1] pqkI *= coulG[p0:p1] # rho_pq(G+k_pq) * conj(rho_sr(G+k_pq)) zdotNC(pqkR, pqkI, rskR.T, rskI.T, 1, eriR, eriI, 1) pqkR = pqkI = rskR = rskI = None return (eriR + eriI * 1j)
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_jk(mydf, dm, hermi=1, kpt=numpy.zeros(3), kpts_band=None, with_j=True, with_k=True, exxdiv=None): '''JK for given k-point''' vj = vk = None if kpts_band is not None and abs(kpt - kpts_band).sum() > 1e-9: kpt = numpy.reshape(kpt, (1, 3)) if with_k: vk = get_k_kpts(mydf, dm, hermi, kpt, kpts_band, exxdiv) if with_j: vj = get_j_kpts(mydf, dm, hermi, kpt, kpts_band) return vj, vk cell = mydf.cell log = logger.Logger(mydf.stdout, mydf.verbose) t1 = (time.clock(), time.time()) dm = numpy.asarray(dm, order='C') dms = _format_dms(dm, [kpt]) nset, _, nao = dms.shape[:3] dms = dms.reshape(nset, nao, nao) j_real = gamma_point(kpt) k_real = gamma_point(kpt) and not numpy.iscomplexobj(dms) kptii = numpy.asarray((kpt, kpt)) kpt_allow = numpy.zeros(3) if with_j: vjcoulG = mydf.weighted_coulG(kpt_allow, False, mydf.gs) vjR = numpy.zeros((nset, nao, nao)) vjI = numpy.zeros((nset, nao, nao)) if with_k: mydf.exxdiv = exxdiv vkcoulG = mydf.weighted_coulG(kpt_allow, True, mydf.gs) vkR = numpy.zeros((nset, nao, nao)) vkI = numpy.zeros((nset, nao, nao)) dmsR = numpy.asarray(dms.real.reshape(nset, nao, nao), order='C') dmsI = numpy.asarray(dms.imag.reshape(nset, nao, nao), order='C') mem_now = lib.current_memory()[0] max_memory = max(2000, (mydf.max_memory - mem_now)) * .8 log.debug1('max_memory = %d MB (%d in use)', max_memory, mem_now) t2 = t1 # rho_rs(-G+k_rs) is computed as conj(rho_{rs^*}(G-k_rs)) # == conj(transpose(rho_sr(G+k_sr), (0,2,1))) blksize = max(int(max_memory * .25e6 / 16 / nao**2), 16) bufR = numpy.empty(blksize * nao**2) bufI = numpy.empty(blksize * nao**2) for pqkR, pqkI, p0, p1 in mydf.pw_loop(mydf.gs, kptii, max_memory=max_memory): t2 = log.timer_debug1('%d:%d ft_aopair' % (p0, p1), *t2) pqkR = pqkR.reshape(nao, nao, -1) pqkI = pqkI.reshape(nao, nao, -1) if with_j: #:v4 = numpy.einsum('ijL,lkL->ijkl', pqk, pqk.conj()) #:vj += numpy.einsum('ijkl,lk->ij', v4, dm) for i in range(nset): rhoR = numpy.einsum('pq,pqk->k', dmsR[i], pqkR) rhoR += numpy.einsum('pq,pqk->k', dmsI[i], pqkI) rhoI = numpy.einsum('pq,pqk->k', dmsI[i], pqkR) rhoI -= numpy.einsum('pq,pqk->k', dmsR[i], pqkI) rhoR *= vjcoulG[p0:p1] rhoI *= vjcoulG[p0:p1] vjR[i] += numpy.einsum('pqk,k->pq', pqkR, rhoR) vjR[i] -= numpy.einsum('pqk,k->pq', pqkI, rhoI) if not j_real: vjI[i] += numpy.einsum('pqk,k->pq', pqkR, rhoI) vjI[i] += numpy.einsum('pqk,k->pq', pqkI, rhoR) #t2 = log.timer_debug1(' with_j', *t2) if with_k: coulG = numpy.sqrt(vkcoulG[p0:p1]) pqkR *= coulG pqkI *= coulG #:v4 = numpy.einsum('ijL,lkL->ijkl', pqk, pqk.conj()) #:vk += numpy.einsum('ijkl,jk->il', v4, dm) pLqR = lib.transpose(pqkR, axes=(0, 2, 1), out=bufR).reshape(-1, nao) pLqI = lib.transpose(pqkI, axes=(0, 2, 1), out=bufI).reshape(-1, nao) iLkR = numpy.ndarray((nao * (p1 - p0), nao), buffer=pqkR) iLkI = numpy.ndarray((nao * (p1 - p0), nao), buffer=pqkI) for i in range(nset): if k_real: lib.dot(pLqR, dmsR[i], 1, iLkR) lib.dot(pLqI, dmsR[i], 1, iLkI) lib.dot(iLkR.reshape(nao, -1), pLqR.reshape(nao, -1).T, 1, vkR[i], 1) lib.dot(iLkI.reshape(nao, -1), pLqI.reshape(nao, -1).T, 1, vkR[i], 1) else: zdotNN(pLqR, pLqI, dmsR[i], dmsI[i], 1, iLkR, iLkI) zdotNC(iLkR.reshape(nao, -1), iLkI.reshape(nao, -1), pLqR.reshape(nao, -1).T, pLqI.reshape(nao, -1).T, 1, vkR[i], vkI[i]) #t2 = log.timer_debug1(' with_k', *t2) pqkR = pqkI = coulG = pLqR = pLqI = iLkR = iLkI = None #t2 = log.timer_debug1('%d:%d'%(p0,p1), *t2) bufR = bufI = None t1 = log.timer_debug1('aft_jk.get_jk', *t1) if with_j: if j_real: vj = vjR else: vj = vjR + vjI * 1j vj = vj.reshape(dm.shape) if with_k: if k_real: vk = vkR else: vk = vkR + vkI * 1j if cell.dimension != 3 and exxdiv: assert (exxdiv.lower() == 'ewald') _ewald_exxdiv_for_G0(cell, kpt, dms, vk) vk = vk.reshape(dm.shape) return vj, vk
def get_k_kpts(mydf, dm_kpts, hermi=1, kpts=numpy.zeros((1, 3)), kpts_band=None, exxdiv=None): cell = mydf.cell log = logger.Logger(mydf.stdout, mydf.verbose) dm_kpts = lib.asarray(dm_kpts, order='C') dms = _format_dms(dm_kpts, kpts) nset, nkpts, nao = dms.shape[:3] swap_2e = (kpts_band is None) kpts_band, input_band = _format_kpts_band(kpts_band, kpts), kpts_band nband = len(kpts_band) kk_table = kpts_band.reshape(-1, 1, 3) - kpts.reshape(1, -1, 3) kk_todo = numpy.ones(kk_table.shape[:2], dtype=bool) vkR = numpy.zeros((nset, nband, nao, nao)) vkI = numpy.zeros((nset, nband, nao, nao)) dmsR = numpy.asarray(dms.real, order='C') dmsI = numpy.asarray(dms.imag, order='C') mem_now = lib.current_memory()[0] max_memory = max(2000, (mydf.max_memory - mem_now)) * .8 log.debug1('max_memory = %d MB (%d in use)', max_memory, mem_now) # K_pq = ( p{k1} i{k2} | i{k2} q{k1} ) def make_kpt(kpt): # kpt = kptj - kpti # search for all possible ki and kj that has ki-kj+kpt=0 kk_match = numpy.einsum('ijx->ij', abs(kk_table + kpt)) < 1e-9 kpti_idx, kptj_idx = numpy.where(kk_todo & kk_match) nkptj = len(kptj_idx) log.debug1('kpt = %s', kpt) log.debug2('kpti_idx = %s', kpti_idx) log.debug2('kptj_idx = %s', kptj_idx) kk_todo[kpti_idx, kptj_idx] = False if swap_2e and not is_zero(kpt): kk_todo[kptj_idx, kpti_idx] = False max_memory1 = max_memory * (nkptj + 1) / (nkptj + 5) #blksize = max(int(max_memory1*4e6/(nkptj+5)/16/nao**2), 16) #bufR = numpy.empty((blksize*nao**2)) #bufI = numpy.empty((blksize*nao**2)) # Use DF object to mimic KRHF/KUHF object in function get_coulG mydf.exxdiv = exxdiv vkcoulG = mydf.weighted_coulG(kpt, True, mydf.gs) kptjs = kpts[kptj_idx] # <r|-G+k_rs|s> = conj(<s|G-k_rs|r>) = conj(<s|G+k_sr|r>) #buf1R = numpy.empty((blksize*nao**2)) #buf1I = numpy.empty((blksize*nao**2)) for aoaoks, p0, p1 in mydf.ft_loop(mydf.gs, kpt, kptjs, max_memory=max_memory1): coulG = numpy.sqrt(vkcoulG[p0:p1]) nG = p1 - p0 bufR = numpy.empty((nG * nao**2)) bufI = numpy.empty((nG * nao**2)) buf1R = numpy.empty((nG * nao**2)) buf1I = numpy.empty((nG * nao**2)) for k, aoao in enumerate(aoaoks): ki = kpti_idx[k] kj = kptj_idx[k] # case 1: k_pq = (pi|iq) #:v4 = numpy.einsum('ijL,lkL->ijkl', pqk, pqk.conj()) #:vk += numpy.einsum('ijkl,jk->il', v4, dm) pLqR = numpy.ndarray((nao, nG, nao), buffer=bufR) pLqI = numpy.ndarray((nao, nG, nao), buffer=bufI) pLqR[:] = aoao.real.reshape(nG, nao, nao).transpose(1, 0, 2) pLqI[:] = aoao.imag.reshape(nG, nao, nao).transpose(1, 0, 2) pLqR *= coulG.reshape(1, nG, 1) pLqI *= coulG.reshape(1, nG, 1) iLkR = numpy.ndarray((nao * nG, nao), buffer=buf1R) iLkI = numpy.ndarray((nao * nG, nao), buffer=buf1I) for i in range(nset): iLkR, iLkI = zdotNN(pLqR.reshape(-1, nao), pLqI.reshape(-1, nao), dmsR[i, kj], dmsI[i, kj], 1, iLkR, iLkI) zdotNC(iLkR.reshape(nao, -1), iLkI.reshape(nao, -1), pLqR.reshape(nao, -1).T, pLqI.reshape(nao, -1).T, 1, vkR[i, ki], vkI[i, ki], 1) # case 2: k_pq = (iq|pi) #:v4 = numpy.einsum('iLj,lLk->ijkl', pqk, pqk.conj()) #:vk += numpy.einsum('ijkl,li->kj', v4, dm) if swap_2e and not is_zero(kpt): iLkR = iLkR.reshape(nao, -1) iLkI = iLkI.reshape(nao, -1) for i in range(nset): iLkR, iLkI = zdotNN(dmsR[i, ki], dmsI[i, ki], pLqR.reshape(nao, -1), pLqI.reshape(nao, -1), 1, iLkR, iLkI) zdotCN( pLqR.reshape(-1, nao).T, pLqI.reshape(-1, nao).T, iLkR.reshape(-1, nao), iLkI.reshape(-1, nao), 1, vkR[i, kj], vkI[i, kj], 1) for ki, kpti in enumerate(kpts_band): for kj, kptj in enumerate(kpts): if kk_todo[ki, kj]: make_kpt(kptj - kpti) if (gamma_point(kpts) and gamma_point(kpts_band) and not numpy.iscomplexobj(dm_kpts)): vk_kpts = vkR else: vk_kpts = vkR + vkI * 1j vk_kpts *= 1. / nkpts # G=0 was not included in the non-uniform grids if cell.dimension != 3 and exxdiv: assert (exxdiv.lower() == 'ewald') _ewald_exxdiv_for_G0(cell, kpts_band, dms, vk_kpts, kpts_band) return _format_jks(vk_kpts, dm_kpts, input_band, kpts)
def get_k_kpts(mydf, dm_kpts, hermi=1, kpts=np.zeros((1, 3)), kpts_band=None, exxdiv=None): '''Get the Coulomb (J) and exchange (K) AO matrices at sampled k-points. Args: dm_kpts : (nkpts, nao, nao) ndarray Density matrix at each k-point kpts : (nkpts, 3) ndarray Kwargs: kpts_band : (3,) ndarray or (*,3) ndarray A list of arbitrary "band" k-points at which to evalute the matrix. Returns: vj : (nkpts, nao, nao) ndarray vk : (nkpts, nao, nao) ndarray or list of vj and vk if the input dm_kpts is a list of DMs ''' cell = mydf.cell gs = mydf.gs coords = cell.gen_uniform_grids(gs) ngs = coords.shape[0] kpts = np.asarray(kpts) dm_kpts = lib.asarray(dm_kpts, order='C') dms = _format_dms(dm_kpts, kpts) nset, nkpts, nao = dms.shape[:3] weight = 1. / nkpts * (cell.vol / ngs) input_band = kpts_band kpts_band, single_kpt_band = _format_kpts_band(kpts_band, kpts) nband = len(kpts_band) if gamma_point(kpts_band) and gamma_point(kpts): vk_kpts = np.zeros((nset, nband, nao, nao), dtype=dms.dtype) else: vk_kpts = np.zeros((nset, nband, nao, nao), dtype=np.complex128) if input_band is None: ao_kpts = mydf._numint.eval_ao(cell, coords, kpts, non0tab=mydf.non0tab) for k2, ao_k2 in enumerate(ao_kpts): kpt2 = kpts[k2] aoR_dms = [lib.dot(ao_k2, dms[i, k2]) for i in range(nset)] for k1, ao_k1 in enumerate(ao_kpts): kpt1 = kpts_band[k1] vkR_k1k2 = get_vkR(mydf, cell, ao_k1, ao_k2, kpt1, kpt2, coords, gs, exxdiv) for i in range(nset): tmp_Rq = np.einsum('Rqs,Rs->Rq', vkR_k1k2, aoR_dms[i]) vk_kpts[i, k1] += weight * lib.dot(ao_k1.T.conj(), tmp_Rq) vkR_k1k2 = aoR_dms = tmp_Rq = None else: for k2, ao_k2 in mydf.aoR_loop(gs, kpts): kpt2 = kpts[k2] aoR_dms = [lib.dot(ao_k2, dms[i, k2]) for i in range(nset)] for k1, ao_k1 in mydf.aoR_loop(gs, kpts_band): kpt1 = kpts_band[k1] vkR_k1k2 = get_vkR(mydf, cell, ao_k1, ao_k2, kpt1, kpt2, coords, gs, exxdiv) for i in range(nset): tmp_Rq = np.einsum('Rqs,Rs->Rq', vkR_k1k2, aoR_dms[i]) vk_kpts[i, k1] += weight * lib.dot(ao_k1.T.conj(), tmp_Rq) vkR_k1k2 = aoR_dms = tmp_Rq = None return _format_jks(vk_kpts, dm_kpts, kpts_band, kpts, single_kpt_band)
def get_ao_pairs_G(mydf, kpts=numpy.zeros((2, 3)), q=None, shls_slice=None, compact=False): '''Calculate forward (G|ij) FFT of all AO pairs. Returns: ao_pairs_G : 2D complex array For gamma point, the shape is (ngs, nao*(nao+1)/2); otherwise the shape is (ngs, nao*nao) ''' if kpts is None: kpts = numpy.zeros((2, 3)) cell = mydf.cell kpts = numpy.asarray(kpts) coords = cell.gen_uniform_grids(mydf.gs) ngs = len(coords) if shls_slice is None: i0, i1 = j0, j1 = (0, cell.nao_nr()) else: ish0, ish1, jsh0, jsh1 = shls_slice ao_loc = cell.ao_loc_nr() i0 = ao_loc[ish0] i1 = ao_loc[ish1] j0 = ao_loc[jsh0] j1 = ao_loc[jsh1] def trans(aoi, aoj, fac=1): if id(aoi) == id(aoj): aoi = aoj = numpy.asarray(aoi.T, order='C') else: aoi = numpy.asarray(aoi.T, order='C') aoj = numpy.asarray(aoj.T, order='C') ni = aoi.shape[0] nj = aoj.shape[0] ao_pairs_G = numpy.empty((ni, nj, ngs), dtype=numpy.complex128) for i in range(ni): ao_pairs_G[i] = tools.fft(fac * aoi[i].conj() * aoj, mydf.gs) ao_pairs_G = ao_pairs_G.reshape(-1, ngs).T return ao_pairs_G if compact and gamma_point(kpts): # gamma point ao = mydf._numint.eval_ao(cell, coords, kpts[:1])[0] ao = numpy.asarray(ao.T, order='C') npair = i1 * (i1 + 1) // 2 - i0 * (i0 + 1) // 2 ao_pairs_G = numpy.empty((npair, ngs), dtype=numpy.complex128) ij = 0 for i in range(i0, i1): ao_pairs_G[ij:ij + i + 1] = tools.fft(ao[i] * ao[:i + 1], mydf.gs) ij += i + 1 ao_pairs_G = ao_pairs_G.T elif is_zero(kpts[0] - kpts[1]): ao = mydf._numint.eval_ao(cell, coords, kpts[:1])[0] ao_pairs_G = trans(ao[:, i0:i1], ao[:, j0:j1]) else: if q is None: q = kpts[1] - kpts[0] aoi, aoj = mydf._numint.eval_ao(cell, coords, kpts[:2]) fac = numpy.exp(-1j * numpy.dot(coords, q)) ao_pairs_G = trans(aoi[:, i0:i1], aoj[:, j0:j1], fac) return ao_pairs_G
def general(mydf, mo_coeffs, kpts=None, compact=False): 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 mo_coeffs = [numpy.asarray(mo, order='F') for mo in mo_coeffs] allreal = not any(numpy.iscomplexobj(mo) for mo in mo_coeffs) q = kptj - kpti coulG = tools.get_coulG(cell, q, gs=mydf.gs) ngs = len(coulG) #################### # gamma point, the integral is real and with s4 symmetry if gamma_point(kptijkl) and allreal: mo_pairs_G = get_mo_pairs_G(mydf, mo_coeffs[:2], kptijkl[:2], q, compact=compact) if ((iden_coeffs(mo_coeffs[0], mo_coeffs[2]) and iden_coeffs(mo_coeffs[1], mo_coeffs[3]))): mo_pairs_G *= numpy.sqrt(coulG).reshape(-1,1) moijR = moklR = mo_pairs_G.real.copy() moijI = moklI = mo_pairs_G.imag.copy() mo_pairs_G = None else: mo_pairs_G *= coulG moijR = mo_pairs_G.real.copy() moijI = mo_pairs_G.imag.copy() mo_pairs_G = None mo_pairs_G = get_mo_pairs_G(mydf, mo_coeffs[2:], kptijkl[2:], q, compact=compact) moklR = mo_pairs_G.real.copy() moklI = mo_pairs_G.imag.copy() mo_pairs_G = None eri = lib.dot(moijR.T, moklR, cell.vol/ngs**2) eri = lib.dot(moijI.T, moklI, cell.vol/ngs**2, eri, 1) return eri #################### # (kpt) i == j == k == l != 0 # (kpt) i == l && j == k && i != j && j != k => # # complex integrals, N^4 elements elif (is_zero(kpti-kptl) and is_zero(kptj-kptk) and iden_coeffs(mo_coeffs[0], mo_coeffs[3]) and iden_coeffs(mo_coeffs[1], mo_coeffs[2])): nmoi = mo_coeffs[0].shape[1] nmoj = mo_coeffs[1].shape[1] mo_ij_G = get_mo_pairs_G(mydf, mo_coeffs[:2], kptijkl[:2]) mo_ij_G *= numpy.sqrt(coulG).reshape(-1,1) mo_kl_G = mo_ij_G.T.reshape(nmoi,nmoj,-1).transpose(1,0,2).conj() mo_kl_G = mo_kl_G.reshape(-1,ngs) return lib.dot(mo_ij_G.T, mo_kl_G.T, cell.vol/ngs**2) #################### # aosym = s1, complex integrals # else: nmok = mo_coeffs[2].shape[1] nmol = mo_coeffs[3].shape[1] mo_ij_G = get_mo_pairs_G(mydf, mo_coeffs[:2], kptijkl[:2], q) mo_ij_G *= coulG.reshape(-1,1) # mo_pairs_invG = rho_rs(-G+k_rs) = conj(rho_sr(G+k_sr)).swap(r,s) mo_kl_G = get_mo_pairs_G(mydf, (mo_coeffs[3],mo_coeffs[2]), (kptl,kptk), q) mo_kl_G = mo_kl_G.T.reshape(nmol,nmok,-1).transpose(1,0,2).conj() mo_kl_G = mo_kl_G.reshape(-1,ngs) return lib.dot(mo_ij_G.T, mo_kl_G.T, cell.vol/ngs**2)
def get_pp(mydf, kpts=None): '''Get the periodic pseudotential nuc-el AO matrix, with G=0 removed. ''' cell = mydf.cell if kpts is None: kpts_lst = numpy.zeros((1, 3)) else: kpts_lst = numpy.reshape(kpts, (-1, 3)) gs = mydf.gs SI = cell.get_SI() Gv = cell.get_Gv(gs) vpplocG = pseudo.get_vlocG(cell, Gv) vpplocG = -numpy.einsum('ij,ij->j', SI, vpplocG) vpplocG[0] = numpy.sum( pseudo.get_alphas(cell)) # from get_jvloc_G0 function ngs = len(vpplocG) # vpploc evaluated in real-space vpplocR = tools.ifft(vpplocG, cell.gs).real vpp = [ lib.dot(aoR.T.conj() * vpplocR, aoR) for k, aoR in mydf.aoR_loop(gs, kpts_lst) ] # vppnonloc evaluated in reciprocal space fakemol = gto.Mole() fakemol._atm = numpy.zeros((1, gto.ATM_SLOTS), dtype=numpy.int32) fakemol._bas = numpy.zeros((1, gto.BAS_SLOTS), dtype=numpy.int32) ptr = gto.PTR_ENV_START fakemol._env = numpy.zeros(ptr + 10) fakemol._bas[0, gto.NPRIM_OF] = 1 fakemol._bas[0, gto.NCTR_OF] = 1 fakemol._bas[0, gto.PTR_EXP] = ptr + 3 fakemol._bas[0, gto.PTR_COEFF] = ptr + 4 # buf for SPG_lmi upto l=0..3 and nl=3 buf = numpy.empty((48, ngs), dtype=numpy.complex128) def vppnl_by_k(kpt): Gk = Gv + kpt G_rad = lib.norm(Gk, axis=1) aokG = ft_ao.ft_ao(cell, Gv, kpt=kpt) * (ngs / cell.vol) vppnl = 0 for ia in range(cell.natm): symb = cell.atom_symbol(ia) if symb not in cell._pseudo: continue pp = cell._pseudo[symb] p1 = 0 for l, proj in enumerate(pp[5:]): rl, nl, hl = proj if nl > 0: fakemol._bas[0, gto.ANG_OF] = l fakemol._env[ptr + 3] = .5 * rl**2 fakemol._env[ptr + 4] = rl**(l + 1.5) * numpy.pi**1.25 pYlm_part = dft.numint.eval_ao(fakemol, Gk, deriv=0) p0, p1 = p1, p1 + nl * (l * 2 + 1) # pYlm is real, SI[ia] is complex pYlm = numpy.ndarray((nl, l * 2 + 1, ngs), dtype=numpy.complex128, buffer=buf[p0:p1]) for k in range(nl): qkl = pseudo.pp._qli(G_rad * rl, l, k) pYlm[k] = pYlm_part.T * qkl #:SPG_lmi = numpy.einsum('g,nmg->nmg', SI[ia].conj(), pYlm) #:SPG_lm_aoG = numpy.einsum('nmg,gp->nmp', SPG_lmi, aokG) #:tmp = numpy.einsum('ij,jmp->imp', hl, SPG_lm_aoG) #:vppnl += numpy.einsum('imp,imq->pq', SPG_lm_aoG.conj(), tmp) if p1 > 0: SPG_lmi = buf[:p1] SPG_lmi *= SI[ia].conj() SPG_lm_aoGs = lib.zdot(SPG_lmi, aokG) p1 = 0 for l, proj in enumerate(pp[5:]): rl, nl, hl = proj if nl > 0: p0, p1 = p1, p1 + nl * (l * 2 + 1) hl = numpy.asarray(hl) SPG_lm_aoG = SPG_lm_aoGs[p0:p1].reshape( nl, l * 2 + 1, -1) tmp = numpy.einsum('ij,jmp->imp', hl, SPG_lm_aoG) vppnl += numpy.einsum('imp,imq->pq', SPG_lm_aoG.conj(), tmp) return vppnl * (1. / ngs**2) for k, kpt in enumerate(kpts_lst): vppnl = vppnl_by_k(kpt) if gamma_point(kpt): vpp[k] = vpp[k].real + vppnl.real else: vpp[k] += vppnl if kpts is None or numpy.shape(kpts) == (3, ): vpp = vpp[0] return numpy.asarray(vpp)
def make_kpt(uniq_kptji_id): # kpt = kptj - kpti kpt = uniq_kpts[uniq_kptji_id] log.debug1('kpt = %s', kpt) adapted_ji_idx = numpy.where(uniq_inverse == uniq_kptji_id)[0] adapted_kptjs = kptjs[adapted_ji_idx] nkptj = len(adapted_kptjs) log.debug1('adapted_ji_idx = %s', adapted_ji_idx) Gaux = ft_ao.ft_ao(fused_cell, Gv, None, b, gxyz, Gvbase, kpt).T Gaux = fuse(Gaux) Gaux *= mydf.weighted_coulG(kpt, False, gs) kLR = Gaux.T.real.copy('C') kLI = Gaux.T.imag.copy('C') j2c = numpy.asarray(feri['j2c/%d' % uniq_kptji_id]) # Note large difference may be found in results between the CD/eig treatments. # In some systems, small integral errors can lead to different treatments of # linear dependency which can be observed in the total energy/orbital energy # around 4th decimal place. # try: # j2c = scipy.linalg.cholesky(j2c, lower=True) # j2ctag = 'CD' # except scipy.linalg.LinAlgError as e: # # Abandon CD treatment for better numerical stablity w, v = scipy.linalg.eigh(j2c) log.debug('MDF metric for kpt %s cond = %.4g, drop %d bfns', uniq_kptji_id, w[0] / w[-1], numpy.count_nonzero(w < df.LINEAR_DEP_THR)) v = v[:, w > df.LINEAR_DEP_THR].T.conj() v /= numpy.sqrt(w[w > df.LINEAR_DEP_THR]).reshape(-1, 1) j2c = v j2ctag = 'eig' naux0 = j2c.shape[0] if is_zero(kpt): # kpti == kptj aosym = 's2' nao_pair = nao * (nao + 1) // 2 vbar = fuse(mydf.auxbar(fused_cell)) ovlp = cell.pbc_intor('int1e_ovlp_sph', hermi=1, kpts=adapted_kptjs) for k, ji in enumerate(adapted_ji_idx): ovlp[k] = lib.pack_tril(ovlp[k]) else: aosym = 's1' nao_pair = nao**2 mem_now = lib.current_memory()[0] log.debug2('memory = %s', mem_now) max_memory = max(2000, mydf.max_memory - mem_now) # nkptj for 3c-coulomb arrays plus 1 Lpq array buflen = min( max(int(max_memory * .6 * 1e6 / 16 / naux / (nkptj + 1)), 1), nao_pair) shranges = _guess_shell_ranges(cell, buflen, aosym) buflen = max([x[2] for x in shranges]) # +1 for a pqkbuf if aosym == 's2': Gblksize = max( 16, int(max_memory * .2 * 1e6 / 16 / buflen / (nkptj + 1))) else: Gblksize = max( 16, int(max_memory * .4 * 1e6 / 16 / buflen / (nkptj + 1))) Gblksize = min(Gblksize, ngs, 16384) pqkRbuf = numpy.empty(buflen * Gblksize) pqkIbuf = numpy.empty(buflen * Gblksize) # buf for ft_aopair buf = numpy.empty((nkptj, buflen * Gblksize), dtype=numpy.complex128) col1 = 0 for istep, sh_range in enumerate(shranges): log.debug1('int3c2e [%d/%d], AO [%d:%d], ncol = %d', \ istep+1, len(shranges), *sh_range) bstart, bend, ncol = sh_range col0, col1 = col1, col1 + ncol j3cR = [] j3cI = [] for k, idx in enumerate(adapted_ji_idx): v = fuse(numpy.asarray(feri['j3c/%d' % idx][:, col0:col1])) if is_zero(kpt): for i, c in enumerate(vbar): if c != 0: v[i] -= c * ovlp[k][col0:col1] j3cR.append(numpy.asarray(v.real, order='C')) if is_zero(kpt) and gamma_point(adapted_kptjs[k]): j3cI.append(None) else: j3cI.append(numpy.asarray(v.imag, order='C')) v = None shls_slice = (bstart, bend, 0, bend) for p0, p1 in lib.prange(0, ngs, Gblksize): dat = ft_ao._ft_aopair_kpts(cell, Gv[p0:p1], shls_slice, aosym, b, gxyz[p0:p1], Gvbase, kpt, adapted_kptjs, out=buf) nG = p1 - p0 for k, ji in enumerate(adapted_ji_idx): aoao = dat[k].reshape(nG, ncol) pqkR = numpy.ndarray((ncol, nG), buffer=pqkRbuf) pqkI = numpy.ndarray((ncol, nG), buffer=pqkIbuf) pqkR[:] = aoao.real.T pqkI[:] = aoao.imag.T lib.dot(kLR[p0:p1].T, pqkR.T, -1, j3cR[k], 1) lib.dot(kLI[p0:p1].T, pqkI.T, -1, j3cR[k], 1) if not (is_zero(kpt) and gamma_point(adapted_kptjs[k])): lib.dot(kLR[p0:p1].T, pqkI.T, -1, j3cI[k], 1) lib.dot(kLI[p0:p1].T, pqkR.T, 1, j3cI[k], 1) for k, ji in enumerate(adapted_ji_idx): if is_zero(kpt) and gamma_point(adapted_kptjs[k]): v = j3cR[k] else: v = j3cR[k] + j3cI[k] * 1j if j2ctag == 'CD': v = scipy.linalg.solve_triangular(j2c, v, lower=True, overwrite_b=True) else: v = lib.dot(j2c, v) feri['j3c/%d' % ji][:naux0, col0:col1] = v del (feri['j2c/%d' % uniq_kptji_id]) for k, ji in enumerate(adapted_ji_idx): v = feri['j3c/%d' % ji][:naux0] del (feri['j3c/%d' % ji]) feri['j3c/%d' % ji] = v
def aux_e2(cell, auxcell, erifile, intor='int3c2e_sph', aosym='s2ij', comp=1, kptij_lst=None, dataname='eri_mo', shls_slice=None, max_memory=2000, verbose=0): r'''3-center AO integrals (ij|L) with double lattice sum: \sum_{lm} (i[l]j[m]|L[0]), where L is the auxiliary basis. On diks, the integrals are stored as (kptij_idx, naux, nao_pair) Args: kptij_lst : (*,2,3) array A list of (kpti, kptj) ''' intor = gto.moleintor.ascint3(intor) if h5py.is_hdf5(erifile): feri = h5py.File(erifile) if dataname in feri: del (feri[dataname]) if dataname + '-kptij' in feri: del (feri[dataname + '-kptij']) else: feri = h5py.File(erifile, 'w') if kptij_lst is None: kptij_lst = numpy.zeros((1, 2, 3)) feri[dataname + '-kptij'] = kptij_lst if shls_slice is None: shls_slice = (0, cell.nbas, 0, cell.nbas, 0, auxcell.nbas) ao_loc = cell.ao_loc_nr() aux_loc = auxcell.ao_loc_nr('ssc' in intor)[:shls_slice[5] + 1] ni = ao_loc[shls_slice[1]] - ao_loc[shls_slice[0]] nj = ao_loc[shls_slice[3]] - ao_loc[shls_slice[2]] naux = aux_loc[shls_slice[5]] - aux_loc[shls_slice[4]] nkptij = len(kptij_lst) nii = (ao_loc[shls_slice[1]] * (ao_loc[shls_slice[1]] + 1) // 2 - ao_loc[shls_slice[0]] * (ao_loc[shls_slice[0]] + 1) // 2) nij = ni * nj kpti = kptij_lst[:, 0] kptj = kptij_lst[:, 1] aosym_ks2 = abs(kpti - kptj).sum(axis=1) < KPT_DIFF_TOL j_only = numpy.all(aosym_ks2) #aosym_ks2 &= (aosym[:2] == 's2' and shls_slice[:2] == shls_slice[2:4]) aosym_ks2 &= aosym[:2] == 's2' for k, kptij in enumerate(kptij_lst): key = '%s/%d' % (dataname, k) if gamma_point(kptij): dtype = 'f8' else: dtype = 'c16' if aosym_ks2[k]: nao_pair = nii else: nao_pair = nij if comp == 1: shape = (naux, nao_pair) else: shape = (comp, naux, nao_pair) chunks = (min(256, naux), min(256, nao_pair)) # 512 KB feri.create_dataset(key, shape, dtype, chunks=chunks) if naux == 0: feri.close() return erifile if j_only and aosym[:2] == 's2': assert (shls_slice[2] == 0) nao_pair = nii else: nao_pair = nij if gamma_point(kptij_lst): dtype = numpy.double else: dtype = numpy.complex128 buflen = max(8, int(max_memory * 1e6 / 16 / (nkptij * ni * nj * comp))) auxdims = aux_loc[shls_slice[4] + 1:shls_slice[5] + 1] - aux_loc[shls_slice[4]:shls_slice[5]] auxranges = balance_segs(auxdims, buflen) buflen = max([x[2] for x in auxranges]) buf = numpy.empty(nkptij * comp * ni * nj * buflen, dtype=dtype) buf1 = numpy.empty(ni * nj * buflen, dtype=dtype) int3c = wrap_int3c(cell, auxcell, intor, aosym, comp, kptij_lst) naux0 = 0 for istep, auxrange in enumerate(auxranges): sh0, sh1, nrow = auxrange sub_slice = (shls_slice[0], shls_slice[1], shls_slice[2], shls_slice[3], shls_slice[4] + sh0, shls_slice[4] + sh1) mat = numpy.ndarray((nkptij, comp, nao_pair, nrow), dtype=dtype, buffer=buf) mat = int3c(sub_slice, mat) for k, kptij in enumerate(kptij_lst): h5dat = feri['%s/%d' % (dataname, k)] for icomp, v in enumerate(mat[k]): v = lib.transpose(v, out=buf1) if gamma_point(kptij): v = v.real if aosym_ks2[k] and v.shape[1] == ni**2: v = lib.pack_tril(v.reshape(-1, ni, ni)) if comp == 1: h5dat[naux0:naux0 + nrow] = v else: h5dat[icomp, naux0:naux0 + nrow] = v naux0 += nrow feri.close() return erifile
def wrap_int3c(cell, auxcell, intor='int3c2e_sph', aosym='s1', comp=1, kptij_lst=numpy.zeros((1, 2, 3))): nbas = cell.nbas atm, bas, env = gto.conc_env(cell._atm, cell._bas, cell._env, cell._atm, cell._bas, cell._env) ao_loc = gto.moleintor.make_loc(bas, intor) aux_loc = auxcell.ao_loc_nr('ssc' in intor) ao_loc = numpy.asarray(numpy.hstack([ao_loc, ao_loc[-1] + aux_loc[1:]]), dtype=numpy.int32) atm, bas, env = gto.conc_env(atm, bas, env, auxcell._atm, auxcell._bas, auxcell._env) Ls = cell.get_lattice_Ls() nimgs = len(Ls) kpti = kptij_lst[:, 0] kptj = kptij_lst[:, 1] if gamma_point(kptij_lst): kk_type = 'g' dtype = numpy.double nkpts = nkptij = 1 kptij_idx = numpy.array([0], dtype=numpy.int32) expkL = numpy.ones(1) elif is_zero(kpti - kptj): # j_only kk_type = 'k' dtype = numpy.complex128 kpts = kptij_idx = kpti expkL = numpy.exp(1j * numpy.dot(kpts, Ls.T)) nkpts = nkptij = len(kpts) else: kk_type = 'kk' dtype = numpy.complex128 kpts = unique(numpy.vstack([kpti, kptj]))[0] expkL = numpy.exp(1j * numpy.dot(kpts, Ls.T)) wherei = numpy.where( abs(kpti.reshape(-1, 1, 3) - kpts).sum(axis=2) < KPT_DIFF_TOL)[1] wherej = numpy.where( abs(kptj.reshape(-1, 1, 3) - kpts).sum(axis=2) < KPT_DIFF_TOL)[1] nkpts = len(kpts) kptij_idx = numpy.asarray(wherei * nkpts + wherej, dtype=numpy.int32) nkptij = len(kptij_lst) fill = 'PBCnr3c_fill_%s%s' % (kk_type, aosym[:2]) drv = libpbc.PBCnr3c_drv cintopt = _vhf.make_cintopt(atm, bas, env, intor) # Remove the precomputed pair data because the pair data corresponds to the # integral of cell #0 while the lattice sum moves shls to all repeated images. libpbc.CINTdel_pairdata_optimizer(cintopt) def int3c(shls_slice, out): shls_slice = (shls_slice[0], shls_slice[1], nbas + shls_slice[2], nbas + shls_slice[3], nbas * 2 + shls_slice[4], nbas * 2 + shls_slice[5]) drv( getattr(libpbc, intor), getattr(libpbc, fill), out.ctypes.data_as(ctypes.c_void_p), ctypes.c_int(nkptij), ctypes.c_int(nkpts), ctypes.c_int(comp), ctypes.c_int(nimgs), Ls.ctypes.data_as(ctypes.c_void_p), expkL.ctypes.data_as(ctypes.c_void_p), kptij_idx.ctypes.data_as(ctypes.c_void_p), (ctypes.c_int * 6)(*shls_slice), ao_loc.ctypes.data_as(ctypes.c_void_p), cintopt, atm.ctypes.data_as(ctypes.c_void_p), ctypes.c_int(cell.natm), bas.ctypes.data_as(ctypes.c_void_p), ctypes.c_int(nbas), # need to pass cell.nbas to libpbc.PBCnr3c_drv env.ctypes.data_as(ctypes.c_void_p)) return out return int3c
def get_k_kpts(mydf, dm_kpts, hermi=1, kpts=numpy.zeros((1, 3)), kpts_band=None, exxdiv=None): cell = mydf.cell log = logger.Logger(mydf.stdout, mydf.verbose) t1 = (time.clock(), time.time()) if mydf._cderi is None or not mydf.has_kpts(kpts_band): if mydf._cderi is not None: log.warn( 'DF integrals for band k-points were not found %s. ' 'DF integrals will be rebuilt to include band k-points.', mydf._cderi) mydf.build(kpts_band=kpts_band) t1 = log.timer_debug1('Init get_k_kpts', *t1) dm_kpts = lib.asarray(dm_kpts, order='C') dms = _format_dms(dm_kpts, kpts) nset, nkpts, nao = dms.shape[:3] kpts_band, input_band = _format_kpts_band(kpts_band, kpts), kpts_band nband = len(kpts_band) vkR = numpy.zeros((nset, nband, nao, nao)) vkI = numpy.zeros((nset, nband, nao, nao)) dmsR = numpy.asarray(dms.real, order='C') dmsI = numpy.asarray(dms.imag, order='C') # K_pq = ( p{k1} i{k2} | i{k2} q{k1} ) bufR = numpy.empty((mydf.blockdim * nao**2)) bufI = numpy.empty((mydf.blockdim * nao**2)) max_memory = max(2000, mydf.max_memory - lib.current_memory()[0]) def make_kpt(ki, kj, swap_2e): kpti = kpts[ki] kptj = kpts_band[kj] for LpqR, LpqI in mydf.sr_loop((kpti, kptj), max_memory, False): nrow = LpqR.shape[0] pLqR = numpy.ndarray((nao, nrow, nao), buffer=bufR) pLqI = numpy.ndarray((nao, nrow, nao), buffer=bufI) tmpR = numpy.ndarray((nao, nrow * nao), buffer=LpqR) tmpI = numpy.ndarray((nao, nrow * nao), buffer=LpqI) pLqR[:] = LpqR.reshape(-1, nao, nao).transpose(1, 0, 2) pLqI[:] = LpqI.reshape(-1, nao, nao).transpose(1, 0, 2) for i in range(nset): zdotNN(dmsR[i, ki], dmsI[i, ki], pLqR.reshape(nao, -1), pLqI.reshape(nao, -1), 1, tmpR, tmpI) zdotCN( pLqR.reshape(-1, nao).T, pLqI.reshape(-1, nao).T, tmpR.reshape(-1, nao), tmpI.reshape(-1, nao), 1, vkR[i, kj], vkI[i, kj], 1) if swap_2e: tmpR = tmpR.reshape(nao * nrow, nao) tmpI = tmpI.reshape(nao * nrow, nao) for i in range(nset): zdotNN(pLqR.reshape(-1, nao), pLqI.reshape(-1, nao), dmsR[i, kj], dmsI[i, kj], 1, tmpR, tmpI) zdotNC(tmpR.reshape(nao, -1), tmpI.reshape(nao, -1), pLqR.reshape(nao, -1).T, pLqI.reshape(nao, -1).T, 1, vkR[i, ki], vkI[i, ki], 1) if kpts_band is None: # normal k-points HF/DFT for ki in range(nkpts): for kj in range(ki): make_kpt(ki, kj, True) make_kpt(ki, ki, False) else: for ki in range(nkpts): for kj in range(nband): make_kpt(ki, kj, False) if (gamma_point(kpts) and gamma_point(kpts_band) and not numpy.iscomplexobj(dm_kpts)): vk_kpts = vkR else: vk_kpts = vkR + vkI * 1j vk_kpts *= 1. / nkpts if exxdiv: assert (exxdiv.lower() == 'ewald') _ewald_exxdiv_for_G0(cell, kpts, dms, vk_kpts, kpts_band) return _format_jks(vk_kpts, dm_kpts, input_band, kpts)
def get_k_kpts(mydf, dm_kpts, hermi=1, kpts=np.zeros((1, 3)), kpts_band=None, exxdiv=None): '''Get the Coulomb (J) and exchange (K) AO matrices at sampled k-points. Args: dm_kpts : (nkpts, nao, nao) ndarray Density matrix at each k-point kpts : (nkpts, 3) ndarray Kwargs: kpts_band : (3,) ndarray or (*,3) ndarray A list of arbitrary "band" k-points at which to evalute the matrix. Returns: vj : (nkpts, nao, nao) ndarray vk : (nkpts, nao, nao) ndarray or list of vj and vk if the input dm_kpts is a list of DMs ''' cell = mydf.cell gs = mydf.gs coords = cell.gen_uniform_grids(gs) ngs = coords.shape[0] if hasattr(dm_kpts, 'mo_coeff'): mo_coeff = dm_kpts.mo_coeff mo_occ = dm_kpts.mo_occ else: mo_coeff = None kpts = np.asarray(kpts) dm_kpts = lib.asarray(dm_kpts, order='C') dms = _format_dms(dm_kpts, kpts) nset, nkpts, nao = dms.shape[:3] weight = 1. / nkpts * (cell.vol / ngs) kpts_band, input_band = _format_kpts_band(kpts_band, kpts), kpts_band nband = len(kpts_band) if gamma_point(kpts_band) and gamma_point(kpts): vk_kpts = np.zeros((nset, nband, nao, nao), dtype=dms.dtype) else: vk_kpts = np.zeros((nset, nband, nao, nao), dtype=np.complex128) ao2_kpts = mydf._numint.eval_ao(cell, coords, kpts, non0tab=mydf.non0tab) ao2_kpts = [np.asarray(ao.T, order='C') for ao in ao2_kpts] if input_band is None: ao1_kpts = ao2_kpts else: ao1_kpts = mydf._numint.eval_ao(cell, coords, kpts_band, non0tab=mydf.non0tab) ao1_kpts = [np.asarray(ao.T, order='C') for ao in ao1_kpts] if mo_coeff is not None and nset == 1: mo_coeff = [ mo_coeff[k][:, occ > 0] * np.sqrt(occ[occ > 0]) for k, occ in enumerate(mo_occ) ] ao2_kpts = [np.dot(mo_coeff[k].T, ao) for k, ao in enumerate(ao2_kpts)] naoj = ao2_kpts[0].shape[0] else: naoj = nao max_memory = mydf.max_memory - lib.current_memory()[0] blksize = int(max(max_memory * 1e6 / 16 / 2 / ngs / nao, 1)) ao1_dtype = np.result_type(*ao1_kpts) ao2_dtype = np.result_type(*ao2_kpts) buf = np.empty((blksize, naoj, ngs), dtype=np.result_type(ao1_dtype, ao2_dtype)) vR_dm = np.empty((nset, nao, ngs), dtype=vk_kpts.dtype) ao_dms = np.empty((nset, naoj, ngs), dtype=np.result_type(dms, ao2_dtype)) for k2, ao2T in enumerate(ao2_kpts): kpt2 = kpts[k2] if mo_coeff is None or nset > 1: for i in range(nset): lib.dot(dms[i, k2], ao2T.conj(), c=ao_dms[i]) else: ao_dms = [ao2T.conj()] for k1, ao1T in enumerate(ao1_kpts): kpt1 = kpts_band[k1] mydf.exxdiv = exxdiv coulG = tools.get_coulG(cell, kpt2 - kpt1, True, mydf, gs) if is_zero(kpt1 - kpt2): expmikr = np.array(1.) else: expmikr = np.exp(-1j * np.dot(coords, kpt2 - kpt1)) for p0, p1 in lib.prange(0, nao, blksize): rho1 = np.einsum('ig,jg->ijg', ao1T[p0:p1].conj() * expmikr, ao2T, out=buf[:p1 - p0]) vG = tools.fft(rho1.reshape(-1, ngs), gs) vG *= coulG vR = tools.ifft(vG, gs).reshape(p1 - p0, naoj, ngs) vG = None if vR_dm.dtype == np.double: vR = vR.real for i in range(nset): np.einsum('ijg,jg->ig', vR, ao_dms[i], out=vR_dm[i, p0:p1]) vR = None vR_dm *= expmikr.conj() for i in range(nset): vk_kpts[i, k1] += weight * lib.dot(vR_dm[i], ao1T.T) return _format_jks(vk_kpts, dm_kpts, input_band, kpts)
def get_j_kpts(mydf, dm_kpts, hermi=1, kpts=numpy.zeros((1, 3)), kpts_band=None): log = logger.Logger(mydf.stdout, mydf.verbose) t1 = (time.clock(), time.time()) if mydf._cderi is None or not mydf.has_kpts(kpts_band): if mydf._cderi is not None: log.warn( 'DF integrals for band k-points were not found %s. ' 'DF integrals will be rebuilt to include band k-points.', mydf._cderi) mydf.build(kpts_band=kpts_band) t1 = log.timer_debug1('Init get_j_kpts', *t1) dm_kpts = lib.asarray(dm_kpts, order='C') dms = _format_dms(dm_kpts, kpts) nset, nkpts, nao = dms.shape[:3] naux = mydf.get_naoaux() nao_pair = nao * (nao + 1) // 2 kpts_band, input_band = _format_kpts_band(kpts_band, kpts), kpts_band nband = len(kpts_band) j_real = gamma_point(kpts_band) and not numpy.iscomplexobj(dms) dmsR = dms.real.transpose(0, 1, 3, 2).reshape(nset, nkpts, nao**2) dmsI = dms.imag.transpose(0, 1, 3, 2).reshape(nset, nkpts, nao**2) rhoR = numpy.zeros((nset, naux)) rhoI = numpy.zeros((nset, naux)) max_memory = max(2000, (mydf.max_memory - lib.current_memory()[0])) for k, kpt in enumerate(kpts): kptii = numpy.asarray((kpt, kpt)) p1 = 0 for LpqR, LpqI in mydf.sr_loop(kptii, max_memory, False): p0, p1 = p1, p1 + LpqR.shape[0] #:Lpq = (LpqR + LpqI*1j).reshape(-1,nao,nao) #:rhoR[:,p0:p1] += numpy.einsum('Lpq,xqp->xL', Lpq, dms[:,k]).real #:rhoI[:,p0:p1] += numpy.einsum('Lpq,xqp->xL', Lpq, dms[:,k]).imag rhoR[:, p0:p1] += numpy.einsum('Lp,xp->xL', LpqR, dmsR[:, k]) rhoI[:, p0:p1] += numpy.einsum('Lp,xp->xL', LpqR, dmsI[:, k]) if LpqI is not None: rhoR[:, p0:p1] -= numpy.einsum('Lp,xp->xL', LpqI, dmsI[:, k]) rhoI[:, p0:p1] += numpy.einsum('Lp,xp->xL', LpqI, dmsR[:, k]) LpqR = LpqI = None t1 = log.timer_debug1('get_j pass 1', *t1) weight = 1. / nkpts rhoR *= weight rhoI *= weight vjR = numpy.zeros((nset, nband, nao_pair)) vjI = numpy.zeros((nset, nband, nao_pair)) for k, kpt in enumerate(kpts_band): kptii = numpy.asarray((kpt, kpt)) p1 = 0 for LpqR, LpqI in mydf.sr_loop(kptii, max_memory, True): p0, p1 = p1, p1 + LpqR.shape[0] #:Lpq = (LpqR + LpqI*1j)#.reshape(-1,nao,nao) #:vjR[:,k] += numpy.dot(rho[:,p0:p1], Lpq).real #:vjI[:,k] += numpy.dot(rho[:,p0:p1], Lpq).imag vjR[:, k] += numpy.dot(rhoR[:, p0:p1], LpqR) if not j_real: vjI[:, k] += numpy.dot(rhoI[:, p0:p1], LpqR) if LpqI is not None: vjR[:, k] -= numpy.dot(rhoI[:, p0:p1], LpqI) vjI[:, k] += numpy.dot(rhoR[:, p0:p1], LpqI) LpqR = LpqI = None t1 = log.timer_debug1('get_j pass 2', *t1) if j_real: vj_kpts = vjR else: vj_kpts = vjR + vjI * 1j vj_kpts = lib.unpack_tril(vj_kpts.reshape(-1, nao_pair)) vj_kpts = vj_kpts.reshape(nset, nband, nao, nao) return _format_jks(vj_kpts, dm_kpts, input_band, kpts)
def make_kpt(uniq_kptji_id): # kpt = kptj - kpti kpt = uniq_kpts[uniq_kptji_id] log.debug1('kpt = %s', kpt) adapted_ji_idx = numpy.where(uniq_inverse == uniq_kptji_id)[0] adapted_kptjs = kptjs[adapted_ji_idx] nkptj = len(adapted_kptjs) log.debug1('adapted_ji_idx = %s', adapted_ji_idx) shls_slice = (auxcell.nbas, fused_cell.nbas) Gaux = ft_ao.ft_ao(fused_cell, Gv, shls_slice, b, gxyz, Gvbase, kpt) Gaux *= mydf.weighted_coulG(kpt, False, gs).reshape(-1, 1) kLR = Gaux.real.copy('C') kLI = Gaux.imag.copy('C') j2c = numpy.asarray(feri['j2c/%d' % uniq_kptji_id]) try: j2c = scipy.linalg.cholesky(j2c, lower=True) j2ctag = 'CD' except scipy.linalg.LinAlgError as e: #msg =('===================================\n' # 'J-metric not positive definite.\n' # 'It is likely that gs is not enough.\n' # '===================================') #log.error(msg) #raise scipy.linalg.LinAlgError('\n'.join([e.message, msg])) w, v = scipy.linalg.eigh(j2c) log.debug('DF metric linear dependency for kpt %s', uniq_kptji_id) log.debug('cond = %.4g, drop %d bfns', w[-1] / w[0], numpy.count_nonzero(w < mydf.linear_dep_threshold)) v = v[:, w > mydf.linear_dep_threshold].T.conj() v /= numpy.sqrt(w[w > mydf.linear_dep_threshold]).reshape(-1, 1) j2c = v j2ctag = 'eig' naux0 = j2c.shape[0] if is_zero(kpt): # kpti == kptj aosym = 's2' nao_pair = nao * (nao + 1) // 2 vbar = fuse(mydf.auxbar(fused_cell)) ovlp = cell.pbc_intor('int1e_ovlp_sph', hermi=1, kpts=adapted_kptjs) for k, ji in enumerate(adapted_ji_idx): ovlp[k] = lib.pack_tril(ovlp[k]) else: aosym = 's1' nao_pair = nao**2 mem_now = lib.current_memory()[0] log.debug2('memory = %s', mem_now) max_memory = max(2000, mydf.max_memory - mem_now) # nkptj for 3c-coulomb arrays plus 1 Lpq array buflen = min( max(int(max_memory * .6 * 1e6 / 16 / naux / (nkptj + 1)), 1), nao_pair) shranges = _guess_shell_ranges(cell, buflen, aosym) buflen = max([x[2] for x in shranges]) # +1 for a pqkbuf if aosym == 's2': Gblksize = max( 16, int(max_memory * .2 * 1e6 / 16 / buflen / (nkptj + 1))) else: Gblksize = max( 16, int(max_memory * .4 * 1e6 / 16 / buflen / (nkptj + 1))) Gblksize = min(Gblksize, ngs, 16384) pqkRbuf = numpy.empty(buflen * Gblksize) pqkIbuf = numpy.empty(buflen * Gblksize) # buf for ft_aopair buf = numpy.empty(nkptj * buflen * Gblksize, dtype=numpy.complex128) col1 = 0 for istep, sh_range in enumerate(shranges): log.debug1('int3c2e [%d/%d], AO [%d:%d], ncol = %d', \ istep+1, len(shranges), *sh_range) bstart, bend, ncol = sh_range col0, col1 = col1, col1 + ncol j3cR = [] j3cI = [] for k, idx in enumerate(adapted_ji_idx): v = numpy.asarray(feri['j3c/%d' % idx][:, col0:col1]) if is_zero(kpt): for i, c in enumerate(vbar): if c != 0: v[i] -= c * ovlp[k][col0:col1] j3cR.append(numpy.asarray(v.real, order='C')) if is_zero(kpt) and gamma_point(adapted_kptjs[k]): j3cI.append(None) else: j3cI.append(numpy.asarray(v.imag, order='C')) v = None shls_slice = (bstart, bend, 0, bend) for p0, p1 in lib.prange(0, ngs, Gblksize): dat = ft_ao._ft_aopair_kpts(cell, Gv[p0:p1], shls_slice, aosym, b, gxyz[p0:p1], Gvbase, kpt, adapted_kptjs, out=buf) nG = p1 - p0 for k, ji in enumerate(adapted_ji_idx): aoao = dat[k].reshape(nG, ncol) pqkR = numpy.ndarray((ncol, nG), buffer=pqkRbuf) pqkI = numpy.ndarray((ncol, nG), buffer=pqkIbuf) pqkR[:] = aoao.real.T pqkI[:] = aoao.imag.T lib.dot(kLR[p0:p1].T, pqkR.T, -1, j3cR[k][naux:], 1) lib.dot(kLI[p0:p1].T, pqkI.T, -1, j3cR[k][naux:], 1) if not (is_zero(kpt) and gamma_point(adapted_kptjs[k])): lib.dot(kLR[p0:p1].T, pqkI.T, -1, j3cI[k][naux:], 1) lib.dot(kLI[p0:p1].T, pqkR.T, 1, j3cI[k][naux:], 1) for k, ji in enumerate(adapted_ji_idx): if is_zero(kpt) and gamma_point(adapted_kptjs[k]): v = fuse(j3cR[k]) else: v = fuse(j3cR[k] + j3cI[k] * 1j) if j2ctag == 'CD': v = scipy.linalg.solve_triangular(j2c, v, lower=True, overwrite_b=True) else: v = lib.dot(j2c, v) feri['j3c/%d' % ji][:naux0, col0:col1] = v del (feri['j2c/%d' % uniq_kptji_id]) for k, ji in enumerate(adapted_ji_idx): v = feri['j3c/%d' % ji][:naux0] del (feri['j3c/%d' % ji]) feri['j3c/%d' % ji] = v