def eval_ao(cell, coords, kpt=None, deriv=0, relativity=0, bastart=0, bascount=None, non0tab=None, verbose=None): '''Collocate AO crystal orbitals (opt. gradients) on the real-space grid. Args: cell : instance of :class:`Cell` coords : (nx*ny*nz, 3) ndarray The real-space grid point coordinates. Kwargs: kpt : (3,) ndarray The k-point corresponding to the crystal AO. deriv : int AO derivative order. It affects the shape of the return array. If deriv=0, the returned AO values are stored in a (N,nao) array. Otherwise the AO values are stored in an array of shape (M,N,nao). Here N is the number of grids, nao is the number of AO functions, M is the size associated to the derivative deriv. Returns: aoR : ([4,] nx*ny*nz, nao=cell.nao_nr()) ndarray The value of the AO crystal orbitals on the real-space grid by default. If deriv=1, also contains the value of the orbitals gradient in the x, y, and z directions. It can be either complex or float array, depending on the kpt argument. If kpt is not given (gamma point), aoR is a float array. See Also: pyscf.dft.numint.eval_ao ''' aoR = 0 for L in tools.get_lattice_Ls(cell, cell.nimgs): if kpt is None: aoR += pyscf.dft.numint.eval_ao(cell, coords-L, deriv, relativity, bastart, bascount, non0tab, verbose) else: factor = numpy.exp(1j*numpy.dot(kpt,L)) aoR += pyscf.dft.numint.eval_ao(cell, coords-L, deriv, relativity, bastart, bascount, non0tab, verbose) * factor if cell.ke_cutoff is not None: ke = 0.5*numpy.einsum('gi,gi->g', cell.Gv, cell.Gv) ke_mask = ke < cell.ke_cutoff aoG = numpy.zeros_like(aoR) for i in range(cell.nao_nr()): if deriv == 1: for c in range(4): aoG[c][ke_mask, i] = tools.fft(aoR[c][:,i], cell.gs)[ke_mask] aoR[c][:,i] = tools.ifft(aoG[c][:,i], cell.gs) else: aoG[ke_mask, i] = tools.fft(aoR[:,i], cell.gs)[ke_mask] aoR[:,i] = tools.ifft(aoG[:,i], cell.gs) return numpy.asarray(aoR)
def _contract_compact(mydf, mos, coulG, max_memory): cell = mydf.cell moiT, mokT = mos nmoi, ngrids = moiT.shape nmok = mokT.shape[0] wcoulG = coulG * (cell.vol/ngrids) def fill_orbital_pair(moT, i0, i1, buf): npair = i1*(i1+1)//2 - i0*(i0+1)//2 out = numpy.ndarray((npair,ngrids), dtype=buf.dtype, buffer=buf) ij = 0 for i in range(i0, i1): numpy.einsum('p,jp->jp', moT[i], moT[:i+1], out=out[ij:ij+i+1]) ij += i + 1 return out eri = numpy.empty((nmoi*(nmoi+1)//2,nmok*(nmok+1)//2)) blksize = int(min(max(nmoi*(nmoi+1)//2, nmok*(nmok+1)//2), (max_memory*1e6/8 - eri.size)/2/ngrids+1)) buf = numpy.empty((blksize,ngrids)) for p0, p1 in lib.prange_tril(0, nmoi, blksize): mo_pairs_G = tools.fft(fill_orbital_pair(moiT, p0, p1, buf), mydf.mesh) mo_pairs_G*= wcoulG v = tools.ifft(mo_pairs_G, mydf.mesh) vR = numpy.asarray(v.real, order='C') for q0, q1 in lib.prange_tril(0, nmok, blksize): mo_pairs = numpy.asarray(fill_orbital_pair(mokT, q0, q1, buf), order='C') eri[p0*(p0+1)//2:p1*(p1+1)//2, q0*(q0+1)//2:q1*(q1+1)//2] = lib.ddot(vR, mo_pairs.T) v = None return eri
def _contract_plain(mydf, mos, coulG, phase, max_memory): cell = mydf.cell moiT, mojT, mokT, molT = mos nmoi, nmoj, nmok, nmol = [x.shape[0] for x in mos] ngrids = moiT.shape[1] wcoulG = coulG * (cell.vol/ngrids) dtype = numpy.result_type(phase, *mos) eri = numpy.empty((nmoi*nmoj,nmok*nmol), dtype=dtype) blksize = int(min(max(nmoi,nmok), (max_memory*1e6/16 - eri.size)/2/ngrids/max(nmoj,nmol)+1)) assert blksize > 0 buf0 = numpy.empty((blksize,max(nmoj,nmol),ngrids), dtype=dtype) buf1 = numpy.ndarray((blksize,nmoj,ngrids), dtype=dtype, buffer=buf0) buf2 = numpy.ndarray((blksize,nmol,ngrids), dtype=dtype, buffer=buf0) for p0, p1 in lib.prange(0, nmoi, blksize): mo_pairs = numpy.einsum('ig,jg->ijg', moiT[p0:p1].conj()*phase, mojT, out=buf1[:p1-p0]) mo_pairs_G = tools.fft(mo_pairs.reshape(-1,ngrids), mydf.mesh) mo_pairs = None mo_pairs_G*= wcoulG v = tools.ifft(mo_pairs_G, mydf.mesh) mo_pairs_G = None v *= phase.conj() if dtype == numpy.double: v = numpy.asarray(v.real, order='C') for q0, q1 in lib.prange(0, nmok, blksize): mo_pairs = numpy.einsum('ig,jg->ijg', mokT[q0:q1].conj(), molT, out=buf2[:q1-q0]) eri[p0*nmoj:p1*nmoj,q0*nmol:q1*nmol] = lib.dot(v, mo_pairs.reshape(-1,ngrids).T) v = None return eri
def test_fft(self): n = 31 a = numpy.random.random([2,n,n,n]) ref = numpy.fft.fftn(a, axes=(1,2,3)).ravel() v = tools.fft(a, [n,n,n]).ravel() self.assertAlmostEqual(abs(ref-v).max(), 0, 10) a = numpy.random.random([2,n,n,8]) ref = numpy.fft.fftn(a, axes=(1,2,3)).ravel() v = tools.fft(a, [n,n,8]).ravel() self.assertAlmostEqual(abs(ref-v).max(), 0, 10) a = numpy.random.random([2,8,n,8]) ref = numpy.fft.fftn(a, axes=(1,2,3)).ravel() v = tools.fft(a, [8,n,8]).ravel() self.assertAlmostEqual(abs(ref-v).max(), 0, 10)
def get_mo_pairs_G(mydf, mo_coeffs, kpts=numpy.zeros((2,3)), q=None, compact=getattr(__config__, 'pbc_df_mo_pairs_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 : (ngrids, 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.mesh) nmoi = mo_coeffs[0].shape[1] nmoj = mo_coeffs[1].shape[1] ngrids = len(coords) def trans(aoi, aoj, fac=1): if id(aoi) == id(aoj) and iden_coeffs(mo_coeffs[0], mo_coeffs[1]): moi = moj = numpy.asarray(lib.dot(mo_coeffs[0].T,aoi.T), order='C') else: moi = numpy.asarray(lib.dot(mo_coeffs[0].T, aoi.T), order='C') moj = numpy.asarray(lib.dot(mo_coeffs[1].T, aoj.T), order='C') mo_pairs_G = numpy.empty((nmoi,nmoj,ngrids), dtype=numpy.complex128) for i in range(nmoi): mo_pairs_G[i] = tools.fft(fac * moi[i].conj() * moj, mydf.mesh) mo_pairs_G = mo_pairs_G.reshape(-1,ngrids).T return mo_pairs_G if gamma_point(kpts): # gamma point, real ao = mydf._numint.eval_ao(cell, coords, kpts[:1])[0] if compact and iden_coeffs(mo_coeffs[0], mo_coeffs[1]): mo = numpy.asarray(lib.dot(mo_coeffs[0].T, ao.T), order='C') npair = nmoi*(nmoi+1)//2 mo_pairs_G = numpy.empty((npair,ngrids), dtype=numpy.complex128) ij = 0 for i in range(nmoi): mo_pairs_G[ij:ij+i+1] = tools.fft(mo[i].conj() * mo[:i+1], mydf.mesh) ij += i + 1 mo_pairs_G = mo_pairs_G.T else: mo_pairs_G = trans(ao, ao) elif is_zero(kpts[0]-kpts[1]): ao = mydf._numint.eval_ao(cell, coords, kpts[:1])[0] mo_pairs_G = trans(ao, ao) else: if q is None: q = kpts[1] - kpts[0] aoi, aoj = mydf._numint.eval_ao(cell, coords, kpts) fac = numpy.exp(-1j * numpy.dot(coords, q)) mo_pairs_G = trans(aoi, aoj, fac) return mo_pairs_G
def get_mo_pairs_G(mydf, mo_coeffs, kpts=numpy.zeros((2,3)), 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 = mydf.kpts cell = mydf.cell kpts = numpy.asarray(kpts) coords = pdft.gen_grid.gen_uniform_grids(cell, 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 ao2mo.incore.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 abs(kpts).sum() < 1e-9: # gamma point, real aoR = mydf._numint.eval_ao(cell, coords, kpts[:1])[0] if compact and ao2mo.incore.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 abs(kpts[0]-kpts[1]).sum() < 1e-9: aoR = mydf._numint.eval_ao(cell, coords, kpts[:1])[0] mo_pairs_G = trans(aoR, aoR) else: aoiR, aojR = mydf._numint.eval_ao(cell, coords, kpts) q = kpts[1] - kpts[0] fac = numpy.exp(-1j * numpy.dot(coords, q)) mo_pairs_G = trans(aoiR, aojR, fac) return mo_pairs_G
def test_ft_ao(self): coords = pdft.gen_grid.gen_uniform_grids(cell) aoR = pdft.numint.eval_ao(cell, coords) ngs, nao = aoR.shape ref = numpy.asarray([tools.fft(aoR[:,i], cell.gs) for i in range(nao)]) ref = ref.T * (cell.vol/ngs) dat = ft_ao.ft_ao(cell, cell.Gv) self.assertAlmostEqual(numpy.linalg.norm(ref[:,0]-dat[:,0]) , 8.4358614794095722e-11, 9) self.assertAlmostEqual(numpy.linalg.norm(ref[:,1]-dat[:,1]) , 0.0041669297531642616 , 4) self.assertAlmostEqual(numpy.linalg.norm(ref[:,2:]-dat[:,2:]), 5.8677286005879366e-14, 9)
def get_vjR(cell, dm, aoR): coulG = tools.get_coulG(cell) rhoR = numint.eval_rho(cell, aoR, dm) rhoG = tools.fft(rhoR, cell.gs) vG = coulG*rhoG vR = tools.ifft(vG, cell.gs) if rhoR.dtype == np.double: vR = vR.real return vR
def trans(aoiR, aojR, fac=1): if id(aoiR) == id(aojR): aoiR = aojR = numpy.asarray(aoiR.T, order='C') else: aoiR = numpy.asarray(aoiR.T, order='C') aojR = numpy.asarray(aojR.T, order='C') ao_pairs_G = numpy.empty((nao,nao,ngs), dtype=numpy.complex128) for i in range(nao): ao_pairs_G[i] = tools.fft(fac * aoiR[i].conj() * aojR, mydf.gs) ao_pairs_G = ao_pairs_G.reshape(-1,ngs).T return ao_pairs_G
def trans(aoiR, aojR, fac=1): if id(aoiR) == id(aojR) and ao2mo.incore.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
def trans(aoi, aoj, fac=1): if id(aoi) == id(aoj) and iden_coeffs(mo_coeffs[0], mo_coeffs[1]): moi = moj = numpy.asarray(lib.dot(mo_coeffs[0].T,aoi.T), order='C') else: moi = numpy.asarray(lib.dot(mo_coeffs[0].T, aoi.T), order='C') moj = numpy.asarray(lib.dot(mo_coeffs[1].T, aoj.T), order='C') mo_pairs_G = numpy.empty((nmoi,nmoj,ngrids), dtype=numpy.complex128) for i in range(nmoi): mo_pairs_G[i] = tools.fft(fac * moi[i].conj() * moj, mydf.mesh) mo_pairs_G = mo_pairs_G.reshape(-1,ngrids).T return mo_pairs_G
def get_j_kpts(mydf, dm_kpts, hermi=1, kpts=np.zeros((1, 3)), kpt_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: kpt_band : (3,) ndarray An arbitrary "band" k-point 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.0 / nkpts rhoG = tools.fft(rhoR[i], gs) vG = coulG * rhoG vR[i] = tools.ifft(vG, gs).real if kpt_band is not None: for k, aoR_kband in mydf.aoR_loop(gs, kpts, kpt_band): pass vj_kpts = [cell.vol / ngs * lib.dot(aoR_kband.T.conj() * vR[i], aoR_kband) for i in range(nset)] if dm_kpts.ndim == 3: # One set of dm_kpts for KRHF vj_kpts = vj_kpts[0] return lib.asarray(vj_kpts) else: vj_kpts = [] weight = cell.vol / ngs for k, aoR in mydf.aoR_loop(gs, kpts): for i in range(nset): vj_kpts.append(weight * lib.dot(aoR.T.conj() * vR[i], aoR)) vj_kpts = lib.asarray(vj_kpts).reshape(nkpts, nset, nao, nao) return vj_kpts.transpose(1, 0, 2, 3).reshape(dm_kpts.shape)
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,ngrids), dtype=numpy.complex128) for i in range(ni): ao_pairs_G[i] = tools.fft(fac * aoi[i].conj() * aoj, mydf.mesh) ao_pairs_G = ao_pairs_G.reshape(-1,ngrids).T return ao_pairs_G
def get_vjR_kpts(cell, dm_kpts, aoR_kpts): nkpts = len(aoR_kpts) coulG = tools.get_coulG(cell) rhoR = 0 for k in range(nkpts): rhoR += 1./nkpts*numint.eval_rho(cell, aoR_kpts[k], dm_kpts[k]) rhoG = tools.fft(rhoR, cell.gs) vG = coulG*rhoG vR = tools.ifft(vG, cell.gs) if rhoR.dtype == np.double: vR = vR.real return vR
def get_ao_pairs_G(mydf, kpts=numpy.zeros((2,3)), 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 = mydf.kpts cell = mydf.cell kpts = numpy.asarray(kpts) coords = pdft.gen_grid.gen_uniform_grids(cell, mydf.gs) nao = cell.nao_nr() ngs = len(coords) def trans(aoiR, aojR, fac=1): if id(aoiR) == id(aojR): aoiR = aojR = numpy.asarray(aoiR.T, order='C') else: aoiR = numpy.asarray(aoiR.T, order='C') aojR = numpy.asarray(aojR.T, order='C') ao_pairs_G = numpy.empty((nao,nao,ngs), dtype=numpy.complex128) for i in range(nao): ao_pairs_G[i] = tools.fft(fac * aoiR[i].conj() * aojR, mydf.gs) ao_pairs_G = ao_pairs_G.reshape(-1,ngs).T return ao_pairs_G if compact and abs(kpts).sum() < 1e-9: # gamma point aoR = mydf._numint.eval_ao(cell, coords, kpts[:1])[0] aoR = numpy.asarray(aoR.T, order='C') npair = nao*(nao+1)//2 ao_pairs_G = numpy.empty((npair,ngs), dtype=numpy.complex128) ij = 0 for i in range(nao): ao_pairs_G[ij:ij+i+1] = tools.fft(aoR[i] * aoR[:i+1], mydf.gs) ij += i + 1 ao_pairs_G = ao_pairs_G.T elif abs(kpts[0]-kpts[1]).sum() < 1e-9: aoR = mydf._numint.eval_ao(cell, coords, kpts[:1])[0] ao_pairs_G = trans(aoR, aoR) else: aoiR, aojR = mydf._numint.eval_ao(cell, coords, kpts[:2]) q = kpts[1] - kpts[0] fac = numpy.exp(-1j * numpy.dot(coords, q)) ao_pairs_G = trans(aoiR, aojR, fac) return ao_pairs_G
def get_ao_pairs_G(cell, kpt=np.zeros(3)): '''Calculate forward (G|ij) and "inverse" (ij|G) FFT of all AO pairs. Args: cell : instance of :class:`Cell` Returns: ao_pairs_G, ao_pairs_invG : (ngs, nao*(nao+1)/2) ndarray The FFTs of the real-space AO pairs. ''' coords = gen_uniform_grids(cell) aoR = eval_ao(cell, coords, kpt) # shape = (coords, nao) ngs, nao = aoR.shape gamma_point = abs(kpt).sum() < 1e-9 if gamma_point: npair = nao*(nao+1)//2 ao_pairs_G = np.empty([ngs, npair], np.complex128) ij = 0 for i in range(nao): for j in range(i+1): ao_ij_R = np.conj(aoR[:,i]) * aoR[:,j] ao_pairs_G[:,ij] = tools.fft(ao_ij_R, cell.gs) #ao_pairs_invG[:,ij] = ngs*tools.ifft(ao_ij_R, cell.gs) ij += 1 ao_pairs_invG = ao_pairs_G.conj() else: ao_pairs_G = np.zeros([ngs, nao,nao], np.complex128) for i in range(nao): for j in range(nao): ao_ij_R = np.conj(aoR[:,i]) * aoR[:,j] ao_pairs_G[:,i,j] = tools.fft(ao_ij_R, cell.gs) ao_pairs_invG = ao_pairs_G.transpose(0,2,1).conj().reshape(-1,nao**2) ao_pairs_G = ao_pairs_G.reshape(-1,nao**2) return ao_pairs_G, ao_pairs_invG
def get_vjR_(cell, aoR, dm): '''Get the real-space Hartree potential of the given density matrix. Returns: vR : (ngs,) ndarray The real-space Hartree potential at every grid point. ''' coulG = tools.get_coulG(cell) rhoR = pyscf.pbc.dft.numint.eval_rho(cell, aoR, dm) rhoG = tools.fft(rhoR, cell.gs) vG = coulG*rhoG vR = tools.ifft(vG, cell.gs) return vR
def test_ft_aoao(self): coords = pdft.gen_grid.gen_uniform_grids(cell) aoR = pdft.numint.eval_ao(cell, coords) ngs, nao = aoR.shape ref = numpy.asarray([tools.fft(aoR[:,i].conj()*aoR[:,j], cell.gs) for i in range(nao) for j in range(nao)]) ref = ref.reshape(nao,nao,-1).transpose(2,0,1) * (cell.vol/ngs) dat = ft_ao.ft_aopair(cell, cell.Gv, aosym='s1hermi') self.assertAlmostEqual(numpy.linalg.norm(ref[:,0,0]-dat[:,0,0]) , 1.869103994619606e-06 , 7) self.assertAlmostEqual(numpy.linalg.norm(ref[:,1,1]-dat[:,1,1]) , 0.02315483195832373 , 4) self.assertAlmostEqual(numpy.linalg.norm(ref[:,2:,2:]-dat[:,2:,2:]), 5.4648896424693173e-14, 9) self.assertAlmostEqual(numpy.linalg.norm(ref[:,0,2:]-dat[:,0,2:]) , 4.0352047774658308e-11, 9) self.assertAlmostEqual(numpy.linalg.norm(ref[:,2:,0]-dat[:,2:,0]) , 4.0352047774658308e-11, 9) idx = numpy.tril_indices(nao) ref = dat[:,idx[0],idx[1]] dat = ft_ao.ft_aopair(cell, cell.Gv, aosym='s2') self.assertAlmostEqual(abs(dat-ref).sum(), 0, 9)
def get_vjR_(cell, dm_kpts, aoR_kpts): '''Get the real-space Hartree potential of the k-point sampled density matrix. Returns: vR : (ngs,) ndarray The real-space Hartree potential at every grid point. ''' nkpts, ngs, nao = aoR_kpts.shape coulG = tools.get_coulG(cell) rhoR = np.zeros(ngs) for k in range(nkpts): rhoR += 1./nkpts*pyscf.pbc.dft.numint.eval_rho(cell, aoR_kpts[k,:,:], dm_kpts[k,:,:]) rhoG = tools.fft(rhoR, cell.gs) vG = coulG*rhoG vR = tools.ifft(vG, cell.gs) return vR
def precompute_exx(self): print "# Precomputing Wigner-Seitz EXX kernel" from pyscf.pbc import gto as pbcgto Nk = tools.get_monkhorst_pack_size(self.cell, self.kpts) print "# Nk =", Nk kcell = pbcgto.Cell() kcell.atom = 'H 0. 0. 0.' kcell.spin = 1 kcell.unit = 'B' kcell.h = self.cell._h * Nk Lc = 1.0/np.linalg.norm(np.linalg.inv(kcell.h.T), axis=0) print "# Lc =", Lc Rin = Lc.min() / 2.0 print "# Rin =", Rin # ASE: alpha = 5./Rin # sqrt(-ln eps) / Rc, eps ~ 10^{-11} kcell.gs = np.array([2*int(L*alpha*3.0) for L in Lc]) # QE: #alpha = 3./Rin * np.sqrt(0.5) #kcell.gs = (4*alpha*np.linalg.norm(kcell.h,axis=0)).astype(int) print "# kcell.gs FFT =", kcell.gs kcell.build(False,False) vR = tools.ifft( tools.get_coulG(kcell), kcell.gs ) kngs = len(vR) print "# kcell kngs =", kngs rs = pyscf.pbc.dft.gen_grid.gen_uniform_grids(kcell) corners = np.dot(np.indices((2,2,2)).reshape((3,8)).T, kcell._h.T) for i, rv in enumerate(rs): # Minimum image convention to corners of kcell parallelepiped r = np.linalg.norm(rv-corners, axis=1).min() if np.isclose(r, 0.): vR[i] = 2*alpha / np.sqrt(np.pi) else: vR[i] = scipy.special.erf(alpha*r) / r vG = (kcell.vol/kngs) * tools.fft(vR, kcell.gs) self.exx_alpha = alpha self.exx_kcell = kcell self.exx_q = kcell.Gv self.exx_vq = vG self.exx_built = True print "# Finished precomputing"
def _fft_quad_integrals(mydf, dm, efg_nuc): # Use FFTDF to compute the integrals of quadrupole operator # (3 \vec{r} \vec{r} - r^2) / r^5 cell = mydf.cell if cell.dimension != 3: raise NotImplementedError mesh = mydf.mesh kpts = mydf.kpts kpts_lst = numpy.reshape(kpts, (-1,3)) nkpts = len(kpts_lst) nao = cell.nao_nr() dm_kpts = dm.reshape((nkpts,nao,nao), order='C') ni = mydf._numint hermi = 1 make_rho, nset, nao = ni._gen_rho_evaluator(cell, dm_kpts, hermi) ngrids = numpy.prod(mesh) rhoR = numpy.zeros(ngrids) for ao_ks_etc, p0, p1 in mydf.aoR_loop(mydf.grids, kpts_lst): ao_ks, mask = ao_ks_etc[0], ao_ks_etc[2] rhoR[p0:p1] += make_rho(0, ao_ks, mask, 'LDA') ao = ao_ks = None rhoG = tools.fft(rhoR, mesh) Gv = cell.get_Gv(mesh) coulG = tools.get_coulG(cell, mesh=mesh, Gv=Gv) GG = numpy.einsum('gx,gy->gxy', Gv, Gv) absG2 = numpy.einsum('gxx->g', GG) # Corresponding to FC term, that makes the tensor traceless idx = numpy.arange(3) GG[:,idx,idx] -= 1./3 * absG2[:,None] vG = 1./ngrids * numpy.einsum('g,g,gxy->gxy', rhoG, coulG, GG) SI = cell.get_SI(Gv) efg_e = lib.einsum('zg,gxy->zxy', SI[efg_nuc], vG.conj()).real return efg_e
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: hermi : int Whether K matrix is hermitian | 0 : not hermitian and not symmetric | 1 : hermitian 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 mesh = mydf.mesh low_dim_ft_type = mydf.low_dim_ft_type coords = cell.gen_uniform_grids(mesh) ngrids = 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 / ngrids) 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) coords = mydf.grids.coords ao2_kpts = [ np.asarray(ao.T, order='C') for ao in mydf._numint.eval_ao(cell, coords, kpts=kpts) ] if input_band is None: ao1_kpts = ao2_kpts else: ao1_kpts = [ np.asarray(ao.T, order='C') for ao in mydf._numint.eval_ao(cell, coords, kpts=kpts_band) ] 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)] mem_now = lib.current_memory()[0] max_memory = mydf.max_memory - mem_now blksize = int( min(nao, max(1, (max_memory - mem_now) * 1e6 / 16 / 4 / ngrids / nao))) lib.logger.debug1(mydf, 'fft_jk: get_k_kpts max_memory %s blksize %d', max_memory, blksize) ao1_dtype = np.result_type(*ao1_kpts) ao2_dtype = np.result_type(*ao2_kpts) vR_dm = np.empty((nset, nao, ngrids), dtype=vk_kpts.dtype) t1 = (time.clock(), time.time()) for k2, ao2T in enumerate(ao2_kpts): if ao2T.size == 0: continue kpt2 = kpts[k2] naoj = ao2T.shape[0] if mo_coeff is None or nset > 1: ao_dms = [lib.dot(dms[i, k2], ao2T.conj()) for i in range(nset)] else: ao_dms = [ao2T.conj()] for k1, ao1T in enumerate(ao1_kpts): kpt1 = kpts_band[k1] # If we have an ewald exxdiv, we add the G=0 correction near the # end of the function to bypass any discretization errors # that arise from the FFT. mydf.exxdiv = exxdiv if exxdiv == 'ewald' or exxdiv is None: coulG = tools.get_coulG(cell, kpt2 - kpt1, False, mydf, mesh, low_dim_ft_type=low_dim_ft_type) else: coulG = tools.get_coulG(cell, kpt2 - kpt1, True, mydf, mesh, low_dim_ft_type=low_dim_ft_type) 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) vG = tools.fft(rho1.reshape(-1, ngrids), mesh) rho1 = None vG *= coulG vR = tools.ifft(vG, mesh).reshape(p1 - p0, naoj, ngrids) 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) t1 = lib.logger.timer_debug1(mydf, 'get_k_kpts: make_kpt (%d,*)' % k2, *t1) # Function _ewald_exxdiv_for_G0 to add back in the G=0 component to vk_kpts # Note in the _ewald_exxdiv_G0 implementation, the G=0 treatments are # different for 1D/2D and 3D systems. The special treatments for 1D and 2D # can only be used with AFTDF/GDF/MDF method. In the FFTDF method, 1D, 2D # and 3D should use the ewald probe charge correction. if exxdiv == 'ewald': _ewald_exxdiv_3d(cell, kpts, dms, vk_kpts, kpts_band=kpts_band) return _format_jks(vk_kpts, dm_kpts, input_band, kpts)
def get_ao_pairs_G(mydf, kpts=numpy.zeros((2,3)), q=None, shls_slice=None, compact=getattr(__config__, 'pbc_df_ao_pairs_compact', False)): '''Calculate forward (G|ij) FFT of all AO pairs. Returns: ao_pairs_G : 2D complex array For gamma point, the shape is (ngrids, nao*(nao+1)/2); otherwise the shape is (ngrids, nao*nao) ''' if kpts is None: kpts = numpy.zeros((2,3)) cell = mydf.cell kpts = numpy.asarray(kpts) coords = cell.gen_uniform_grids(mydf.mesh) ngrids = 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,ngrids), dtype=numpy.complex128) for i in range(ni): ao_pairs_G[i] = tools.fft(fac * aoi[i].conj() * aoj, mydf.mesh) ao_pairs_G = ao_pairs_G.reshape(-1,ngrids).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,ngrids), 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.mesh) 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 get_k_kpts_occ(mydf, dm_kpts, hermi=1, kpts=numpy.zeros((1, 3)), kpts_band=None, exxdiv=None): mydf = _sync_mydf(mydf) cell = mydf.cell mesh = mydf.mesh coords = cell.gen_uniform_grids(mesh) ngrids = coords.shape[0] if hasattr(dm_kpts, 'mo_coeff'): if dm_kpts.ndim == 3: # KRHF mo_coeff = [dm_kpts.mo_coeff] mo_occ = [dm_kpts.mo_occ] else: # KUHF mo_coeff = dm_kpts.mo_coeff mo_occ = dm_kpts.mo_occ elif hasattr(dm_kpts[0], 'mo_coeff'): mo_coeff = [dm.mo_coeff for dm in dm_kpts] mo_occ = [dm.mo_occ for dm in dm_kpts] else: mo_coeff = None kpts = numpy.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 / ngrids) 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 = numpy.zeros((nset, nband, nao, nao), dtype=dms.dtype) else: vk_kpts = numpy.zeros((nset, nband, nao, nao), dtype=numpy.complex128) coords = mydf.grids.coords ao2_kpts = [ numpy.asarray(ao.T, order='C') for ao in mydf._numint.eval_ao(cell, coords, kpts=kpts) ] if input_band is None: ao1_kpts = ao2_kpts else: ao1_kpts = [ numpy.asarray(ao.T, order='C') for ao in mydf._numint.eval_ao(cell, coords, kpts=kpts_band) ] # occ if mo_coeff is None: return _format_jks(vk_kpts, dm_kpts, input_band, kpts) # occ if gamma_point(kpts_band) and gamma_point(kpts): for i in range(nset): occ = mo_occ[i][0] kiv = numpy.zeros( (nset, nband, mo_coeff[i][0][:, occ > 0].shape[1], nao), dtype=dms.dtype) else: kiv = [[] * nset] for i in range(nset): for k1 in range(nband): occ = mo_occ[i][k1] kiv[i].append( numpy.zeros((mo_coeff[i][k1][:, occ > 0].shape[1], nao), dtype=numpy.complex128)) mem_now = lib.current_memory()[0] max_memory = mydf.max_memory - mem_now blksize = int( min(nao, max(1, (max_memory - mem_now) * 1e6 / 16 / 4 / ngrids / nao))) lib.logger.debug1(mydf, 'max_memory %s blksize %d', max_memory, blksize) ao1_dtype = numpy.result_type(*ao1_kpts) ao2_dtype = numpy.result_type(*ao2_kpts) vR_dm = numpy.empty((nset, nao, ngrids), dtype=vk_kpts.dtype) ao_dms_buf = [None] * nkpts # occ ao3T_buf = [None] * nkpts tasks = [(k1, k2) for k2 in range(nkpts) for k1 in range(nband)] for k1, k2 in mpi.static_partition(tasks): ao1T = ao1_kpts[k1] ao2T = ao2_kpts[k2] kpt1 = kpts_band[k1] kpt2 = kpts[k2] if ao2T.size == 0 or ao1T.size == 0: continue # If we have an ewald exxdiv, we add the G=0 correction near the # end of the function to bypass any discretization errors # that arise from the FFT. mydf.exxdiv = exxdiv if exxdiv == 'ewald' or exxdiv is None: coulG = tools.get_coulG(cell, kpt2 - kpt1, False, mydf, mesh) else: coulG = tools.get_coulG(cell, kpt2 - kpt1, True, mydf, mesh) if is_zero(kpt1 - kpt2): expmikr = numpy.array(1.) else: expmikr = numpy.exp(-1j * numpy.dot(coords, kpt2 - kpt1)) if ao_dms_buf[k2] is None: if mo_coeff is None: ao_dms = [lib.dot(dm[k2], ao2T.conj()) for dm in dms] else: ao_dms = [] for i, dm in enumerate(dms): occ = mo_occ[i][k2] mo_scaled = mo_coeff[i][k2][:, occ > 0] * numpy.sqrt( occ[occ > 0]) ao_dms.append(lib.dot(mo_scaled.T, ao2T).conj()) ao_dms_buf[k2] = ao_dms else: ao_dms = ao_dms_buf[k2] if ao3T_buf[k1] is None: if mo_coeff is not None: ao3T = [] for i, dm in enumerate(dms): occ = mo_occ[i][k1] mo_scaled = mo_coeff[i][k1][:, occ > 0] ao3T.append(lib.dot(mo_scaled.T, ao1T)) ao3T_buf[k1] = ao3T else: ao3T = ao3T_buf[k1] if mo_coeff is None: for p0, p1 in lib.prange(0, nao, blksize): rho1 = numpy.einsum('ig,jg->ijg', ao1T[p0:p1].conj() * expmikr, ao2T) vG = tools.fft(rho1.reshape(-1, ngrids), mesh) rho1 = None vG *= coulG vR = tools.ifft(vG, mesh).reshape(p1 - p0, nao, ngrids) vG = None if vR_dm.dtype == numpy.double: vR = vR.real for i in range(nset): numpy.einsum('ijg,jg->ig', vR, ao_dms[i], out=vR_dm[i, p0:p1]) vR = None else: for i in range(nset): for p0, p1 in lib.prange(0, ao3T[i].shape[0], blksize): rho1 = numpy.einsum('ig,jg->ijg', ao3T[i][p0:p1].conj() * expmikr, ao_dms[i].conj()) vG = tools.fft(rho1.reshape(-1, ngrids), mesh) rho1 = None vG *= coulG vR = tools.ifft(vG, mesh).reshape(p1 - p0, -1, ngrids) vG = None if vR_dm.dtype == numpy.double: vR = vR.real numpy.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): kiv[i][k1] += weight * lib.dot( vR_dm[i, 0:mo_coeff[i][k1][:, mo_occ[i][k1] > 0].shape[1]], ao1T.T) kiv = mpi.reduce(lib.asarray(kiv)) if rank == 0: for i in range(nset): for k1 in range(nband): kij = lib.einsum('ui,ju->ij', mo_coeff[i][k1][:, mo_occ[i][k1] > 0], kiv[i][k1]) kr = scipy.linalg.solve(kij.conj(), kiv[i][k1]) vk_kpts[i, k1] = lib.dot(kiv[i][k1].T.conj(), kr) if gamma_point(kpts_band) and gamma_point(kpts): vk_kpts = vk_kpts.real if rank == 0: if exxdiv == 'ewald': _ewald_exxdiv_for_G0(cell, kpts, dms, vk_kpts, kpts_band=kpts_band) return _format_jks(vk_kpts, dm_kpts, input_band, kpts)
cell.a = numpy.diag([L, L, L]) cell.mesh = numpy.array([n, n, n]) cell.atom = '''C 1.3 .2 .3 C .1 .1 1.1 ''' cell.basis = 'ccpvdz' #cell.basis = {'C': [[0, (2.4, .1, .6), (1.0,.8, .4)], [1, (1.1, 1)]]} #cell.basis = {'C': [[0, (2.4, 1)]]} cell.unit = 'B' #cell.verbose = 4 cell.build(0, 0) #cell.nimgs = (2,2,2) ao2 = ft_aopair(cell, cell.Gv) nao = cell.nao_nr() coords = pyscf.pbc.dft.gen_grid.gen_uniform_grids(cell) aoR = cell.pbc_eval_gto('GTOval', coords) aoR2 = numpy.einsum('ki,kj->kij', aoR.conj(), aoR) ngrids = aoR.shape[0] for i in range(nao): for j in range(nao): ao2ref = tools.fft(aoR2[:, i, j], cell.mesh) * cell.vol / ngrids print(i, j, numpy.linalg.norm(ao2ref - ao2[:, i, j])) aoG = ft_ao(cell, cell.Gv) for i in range(nao): aoref = tools.fft(aoR[:, i], cell.mesh) * cell.vol / ngrids print(i, numpy.linalg.norm(aoref - aoG[:, i]))
def get_j_e1_kpts(mydf, dm_kpts, kpts=np.zeros((1, 3)), kpts_band=None): '''Derivatives of Coulomb (J) AO matrix at sampled k-points. ''' cell = mydf.cell mesh = mydf.mesh ni = mydf._numint make_rho, nset, nao = ni._gen_rho_evaluator(cell, dm_kpts, hermi=1) 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, mesh=mesh) ngrids = len(coulG) if gamma_point(kpts): vR = rhoR = np.zeros((nset, ngrids)) for ao_ks_etc, p0, p1 in mydf.aoR_loop(mydf.grids, kpts): ao_ks, mask = ao_ks_etc[0], ao_ks_etc[2] for i in range(nset): rhoR[i, p0:p1] += make_rho(i, ao_ks, mask, 'LDA') ao = ao_ks = None for i in range(nset): rhoG = tools.fft(rhoR[i], mesh) vG = coulG * rhoG vR[i] = tools.ifft(vG, mesh).real else: # vR may be complex if the underlying density is complex vR = rhoR = np.zeros((nset, ngrids), dtype=np.complex128) for ao_ks_etc, p0, p1 in mydf.aoR_loop(mydf.grids, kpts): ao_ks, mask = ao_ks_etc[0], ao_ks_etc[2] for i in range(nset): for k, ao in enumerate(ao_ks): ao_dm = lib.dot(ao, dms[i, k]) rhoR[i, p0:p1] += np.einsum('xi,xi->x', ao_dm, ao.conj()) rhoR *= 1. / nkpts for i in range(nset): rhoG = tools.fft(rhoR[i], mesh) vG = coulG * rhoG vR[i] = tools.ifft(vG, mesh) kpts_band, input_band = _format_kpts_band(kpts_band, kpts), kpts_band nband = len(kpts_band) weight = cell.vol / ngrids vR *= weight if gamma_point(kpts_band): vj_kpts = np.zeros((3, nset, nband, nao, nao)) else: vj_kpts = np.zeros((3, nset, nband, nao, nao), dtype=np.complex128) for ao_ks_etc, p0, p1 in mydf.aoR_loop(mydf.grids, kpts_band, deriv=1): ao_ks, mask = ao_ks_etc[0], ao_ks_etc[2] for i in range(nset): # ni.eval_mat can handle real vR only # vj_kpts[i] += ni.eval_mat(cell, ao_ks, 1., None, vR[i,p0:p1], mask, 'LDA') for k, ao in enumerate(ao_ks): aow = np.einsum('xi,x->xi', ao[0], vR[i, p0:p1]) vj_kpts[:, i, k] -= lib.einsum('axi,xj->aij', ao[1:].conj(), aow) vj_kpts = np.asarray( [_format_jks(vj, dm_kpts, input_band, kpts) for vj in vj_kpts]) return vj_kpts
def get_k_e1_kpts(mydf, dm_kpts, kpts=np.zeros((1, 3)), kpts_band=None, exxdiv=None): '''Derivatives of exchange (K) AO matrix at sampled k-points. ''' cell = mydf.cell mesh = mydf.mesh coords = cell.gen_uniform_grids(mesh) ngrids = coords.shape[0] if getattr(dm_kpts, 'mo_coeff', None) is not None: 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 / ngrids) 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((3, nset, nband, nao, nao), dtype=dms.dtype) else: vk_kpts = np.zeros((3, nset, nband, nao, nao), dtype=np.complex128) coords = mydf.grids.coords if input_band is None: ao2_kpts = [ np.asarray(ao.transpose(0, 2, 1), order='C') for ao in mydf._numint.eval_ao(cell, coords, kpts=kpts, deriv=1) ] ao1_kpts = ao2_kpts ao2_kpts = [ao2_kpt[0] for ao2_kpt in ao2_kpts] else: ao2_kpts = [ np.asarray(ao.T, order='C') for ao in mydf._numint.eval_ao(cell, coords, kpts=kpts) ] ao1_kpts = [ np.asarray(ao.transpose(0, 2, 1), order='C') for ao in mydf._numint.eval_ao(cell, coords, kpts=kpts_band, deriv=1) ] 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)] mem_now = lib.current_memory()[0] max_memory = mydf.max_memory - mem_now blksize = int( min(nao, max(1, (max_memory - mem_now) * 1e6 / 16 / 4 / 3 / ngrids / nao))) lib.logger.debug1(mydf, 'fft_jk: get_k_kpts max_memory %s blksize %d', max_memory, blksize) vR_dm = np.empty((3, nset, nao, ngrids), dtype=vk_kpts.dtype) t1 = (time.clock(), time.time()) for k2, ao2T in enumerate(ao2_kpts): if ao2T.size == 0: continue kpt2 = kpts[k2] naoj = ao2T.shape[0] if mo_coeff is None or nset > 1: ao_dms = [lib.dot(dms[i, k2], ao2T.conj()) for i in range(nset)] else: ao_dms = [ao2T.conj()] for k1, ao1T in enumerate(ao1_kpts): kpt1 = kpts_band[k1] # If we have an ewald exxdiv, we add the G=0 correction near the # end of the function to bypass any discretization errors # that arise from the FFT. mydf.exxdiv = exxdiv if exxdiv == 'ewald' or exxdiv is None: coulG = tools.get_coulG(cell, kpt2 - kpt1, False, mydf, mesh) else: coulG = tools.get_coulG(cell, kpt2 - kpt1, True, mydf, mesh) 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('aig,jg->aijg', ao1T[1:, p0:p1].conj() * expmikr, ao2T) vG = tools.fft(rho1.reshape(-1, ngrids), mesh) rho1 = None vG *= coulG vR = tools.ifft(vG, mesh).reshape(3, p1 - p0, naoj, ngrids) vG = None if vR_dm.dtype == np.double: vR = vR.real for i in range(nset): np.einsum('aijg,jg->aig', 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 * np.einsum( 'aig,jg->aij', vR_dm[:, i], ao1T[0]) t1 = lib.logger.timer_debug1(mydf, 'get_k_kpts: make_kpt (%d,*)' % k2, *t1) # Ewald correction has no contribution to nuclear gradient unless range separted Coulomb is used # The gradient correction part is not added in the vk matrix if exxdiv == 'ewald' and cell.omega != 0: raise NotImplementedError("Range Separated Coulomb") # when cell.omega !=0: madelung constant will have a non-zero derivative vk_kpts = np.asarray( [_format_jks(vk, dm_kpts, input_band, kpts) for vk in vk_kpts]) return vk_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 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) mo_ids = [id(x) for x in mo_coeff_kpts] moTs = [] coords = cell.gen_uniform_grids(mydf.mesh) aos = mydf._numint.eval_ao(cell, coords, kpts) for n, mo_id in enumerate(mo_ids): if mo_id in mo_ids[:n]: moTs.append(moTs[mo_ids[:n].index(mo_id)]) else: moTs.append([lib.dot(mo.T, aos[k].T) for k,mo in enumerate(mo_coeff_kpts[n])]) # 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) # To hold intermediates fswap = lib.H5TmpFile() 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] ki = adapted_ji_idx[0] // nkpts kj = adapted_ji_idx[0] % nkpts coulG = tools.get_coulG(cell, q, mesh=mydf.mesh) coulG *= (cell.vol/ngrids) * factor phase = numpy.exp(-1j * numpy.dot(coords, q)) for kk in range(nkpts): kl = kconserv[ki, kj, kk] mokT = moTs[2][kk] molT = moTs[3][kl] mo_pairs = numpy.einsum('ig,g,jg->ijg', mokT.conj(), phase.conj(), molT) v = tools.ifft(mo_pairs.reshape(-1,ngrids), mydf.mesh) v *= coulG v = tools.fft(v.reshape(-1,ngrids), mydf.mesh) v *= phase fswap['zkl/'+str(kk)] = v for ji_idx in adapted_ji_idx: ki = ji_idx // nkpts kj = ji_idx % nkpts for kk in range(nkpts): moiT = moTs[0][ki] mojT = moTs[1][kj] mo_pairs = numpy.einsum('ig,jg->ijg', moiT.conj(), mojT) tmp = lib.dot(mo_pairs.reshape(-1,ngrids), numpy.asarray(fswap['zkl/'+str(kk)]).T) if dtype == numpy.double: tmp = tmp.real out[ki,kj,kk] = tmp.reshape(eri_shape[3:]) del(fswap['zkl']) return out
def uks_j_xc(mydf, dm_kpts, xc_code, hermi=1, kpts=numpy.zeros((1, 3)), kpts_band=None, with_j=WITH_J, j_in_xc=J_IN_XC): log = lib.logger.Logger(mydf.stdout, mydf.verbose) cell = mydf.cell dm_kpts = lib.asarray(dm_kpts, order='C') dms = _format_dms(dm_kpts, kpts) nset, nkpts, nao = dms.shape[:3] dms = None #TODO: Handle multiple sets of KUKS density matrices (2,nset,nkpts,nao,nao) assert (nset == 2) # alpha and beta density matrices in KUKS ni = mydf._numint xctype = ni._xc_type(xc_code) if xctype == 'LDA': deriv = 0 rhoG = _eval_rhoG(mydf, dm_kpts, hermi, kpts, deriv=0) def add_j_(v, ao_l, ao_h, idx_l, idx_h, vR): for k in range(nkpts): aow = numpy.einsum('pi,p->pi', ao_l[k], vR[0]) v[0, k, idx_l[:, None], idx_h] += lib.dot(aow.conj().T, ao_h[k]) aow = numpy.einsum('pi,p->pi', ao_l[k], vR[1]) v[1, k, idx_l[:, None], idx_h] += lib.dot(aow.conj().T, ao_h[k]) def add_xc_(v, ao_l, ao_h, idx_l, idx_h, wv): add_j_(v, ao_l, ao_h, idx_l, idx_h, wv[:, 0]) elif xctype == 'GGA': deriv = 1 if RHOG_HIGH_DERIV: rhoG = _eval_rhoG(mydf, dm_kpts, hermi, kpts, deriv) else: Gv = cell.Gv ngrids = Gv.shape[0] rhoG = numpy.empty((2, 4, ngrids), dtype=numpy.complex128) rhoG[:, :1] = _eval_rhoG(mydf, dm_kpts, hermi, kpts, deriv=0) rhoG[:, 1:] = numpy.einsum('np,px->nxp', 1j * rhoG[:, 0], Gv) def add_j_(v, ao_l, ao_h, idx_l, idx_h, vR): for k in range(nkpts): aow = numpy.einsum('pi,p->pi', ao_l[k][0], vR[0]) v[0, k, idx_l[:, None], idx_h] += lib.dot(aow.conj().T, ao_h[k][0]) aow = numpy.einsum('pi,p->pi', ao_l[k][0], vR[1]) v[1, k, idx_l[:, None], idx_h] += lib.dot(aow.conj().T, ao_h[k][0]) def add_xc_(v, ao_l, ao_h, idx_l, idx_h, wv): wva, wvb = wv for k in range(nkpts): aow = numpy.einsum('npi,np->pi', ao_l[k][:4], wva) v1 = lib.dot(aow.conj().T, ao_h[k][0]) aow = numpy.einsum('npi,np->pi', ao_h[k][1:4], wva[1:4]) v1 += lib.dot(ao_l[k][0].conj().T, aow) v[0, k, idx_l[:, None], idx_h] += v1 aow = numpy.einsum('npi,np->pi', ao_l[k][:4], wvb) v1 = lib.dot(aow.conj().T, ao_h[k][0]) aow = numpy.einsum('npi,np->pi', ao_h[k][1:4], wvb[1:4]) v1 += lib.dot(ao_l[k][0].conj().T, aow) v[1, k, idx_l[:, None], idx_h] += v1 else: # MGGA deriv = 2 #TODO: RHOG_HIGH_DERIV: rhoG = _eval_rhoG(mydf, dm_kpts, hermi, kpts, deriv) def add_j_(v, ao_l, ao_h, idx_l, idx_h, vR): for k in range(nkpts): aow = numpy.einsum('pi,p->pi', ao_l[k][0], vR[0]) v[0, k, idx_l[:, None], idx_h] += lib.dot(aow.conj().T, ao_h[k][0]) aow = numpy.einsum('pi,p->pi', ao_l[k][0], vR[1]) v[1, k, idx_l[:, None], idx_h] += lib.dot(aow.conj().T, ao_h[k][0]) def add_xc_(v, ao_l, ao_h, idx_l, idx_h, wv): wva, wvb = wv for k in range(nkpts): aow = numpy.einsum('npi,np->pi', ao_l[k][:4], wva[:4]) v1 = lib.dot(aow.conj().T, ao_h[k][0]) aow = numpy.einsum('npi,np->pi', ao_h[k][1:4], wva[1:4]) v1 += lib.dot(ao_l[k][0].conj().T, aow) aow = numpy.einsum('pi,p->pi', ao_h[k][1], wva[4], out=aow) v1 += lib.dot(ao_l[k][1].conj().T, aow) aow = numpy.einsum('pi,p->pi', ao_h[k][2], wva[4], out=aow) v1 += lib.dot(ao_l[k][2].conj().T, aow) aow = numpy.einsum('pi,p->pi', ao_h[k][3], wva[4], out=aow) v1 += lib.dot(ao_l[k][3].conj().T, aow) v[0, k, idx_l[:, None], idx_h] += v1 aow = numpy.einsum('npi,np->pi', ao_l[k][:4], wvb[:4]) v1 = lib.dot(aow.conj().T, ao_h[k][0]) aow = numpy.einsum('npi,np->pi', ao_h[k][1:4], wvb[1:4]) v1 += lib.dot(ao_l[k][0].conj().T, aow) aow = numpy.einsum('pi,p->pi', ao_h[k][1], wvb[4], out=aow) v1 += lib.dot(ao_l[k][1].conj().T, aow) aow = numpy.einsum('pi,p->pi', ao_h[k][2], wvb[4], out=aow) v1 += lib.dot(ao_l[k][2].conj().T, aow) aow = numpy.einsum('pi,p->pi', ao_h[k][3], wvb[4], out=aow) v1 += lib.dot(ao_l[k][3].conj().T, aow) v[1, k, idx_l[:, None], idx_h] += v1 mesh = cell.mesh coulG = tools.get_coulG(cell, mesh=mesh, low_dim_ft_type=mydf.low_dim_ft_type) ngrids = coulG.size vG = numpy.einsum('ng,g->ng', rhoG[:, 0].reshape(-1, ngrids), coulG) vG = vG.reshape(2, *mesh) weight = cell.vol / ngrids # *(1./weight) because rhoR is scaled by weight in _eval_rhoG. When # computing rhoR with IFFT, the weight factor is not needed. rhoR = tools.ifft(rhoG.reshape(-1, ngrids), mesh) * (1. / weight) rhoR = rhoR.real.reshape(2, -1, ngrids) nelec = numpy.zeros(2) excsum = 0 exc, vxc = ni.eval_xc(xc_code, rhoR, 1, deriv=1)[:2] if xctype == 'LDA': vrho = vxc[0] wva = vrho[:, 0].reshape(1, ngrids) wvb = vrho[:, 1].reshape(1, ngrids) elif xctype == 'GGA': vrho, vsigma = vxc[:2] wva = numpy.empty((4, ngrids)) wvb = numpy.empty((4, ngrids)) wva[0] = vrho[:, 0] wva[1:4] = rhoR[0, 1:4] * (vsigma[:, 0] * 2) # sigma_uu wva[1:4] += rhoR[1, 1:4] * vsigma[:, 1] # sigma_ud wvb[0] = vrho[:, 1] wvb[1:4] = rhoR[1, 1:4] * (vsigma[:, 2] * 2) # sigma_dd wvb[1:4] += rhoR[0, 1:4] * vsigma[:, 1] # sigma_ud else: vrho, vsigma, vlapl, vtau = vxc wva = numpy.empty((5, ngrids)) wvb = numpy.empty((5, ngrids)) wva[0] = vrho[:, 0] wva[1:4] = rhoR[0, 1:4] * (vsigma[:, 0] * 2) # sigma_uu wva[1:4] += rhoR[1, 1:4] * vsigma[:, 1] # sigma_ud wvb[0] = vrho[:, 1] wvb[1:4] = rhoR[1, 1:4] * (vsigma[:, 2] * 2) # sigma_dd wvb[1:4] += rhoR[0, 1:4] * vsigma[:, 1] # sigma_ud if vlapl is None: wvb[4] = .5 * vtau[:, 1] wva[4] = .5 * vtau[:, 0] else: wva[4] = (.5 * vtau[:, 0] + 2 * vlapl[:, 0]) wvb[4] = (.5 * vtau[:, 1] + 2 * vlapl[:, 1]) nelec[0] += rhoR[0, 0].sum() * weight nelec[1] += rhoR[1, 0].sum() * weight excsum += (rhoR[0, 0] * exc).sum() * weight excsum += (rhoR[1, 0] * exc).sum() * weight wv_freq = tools.fft(numpy.vstack((wva, wvb)), mesh) * weight wv_freq = wv_freq.reshape(2, -1, *mesh) if j_in_xc: wv_freq[:, 0] += vG vR = tools.ifft(vG.reshape(-1, ngrids), mesh) ecoul = numpy.einsum('ng,ng->', rhoR[:, 0].real, vR.real) * .5 log.debug('Coulomb energy %s', ecoul) excsum += ecoul rhoR = rhoG = None kpts_band, input_band = _format_kpts_band(kpts_band, kpts), kpts_band if gamma_point(kpts_band): veff = numpy.zeros((2, nkpts, nao, nao)) vj = numpy.zeros((2, nkpts, nao, nao)) else: veff = numpy.zeros((2, nkpts, nao, nao), dtype=numpy.complex128) vj = numpy.zeros((2, nkpts, nao, nao), dtype=numpy.complex128) for grids_high, grids_low in mydf.tasks: cell_high = grids_high.cell mesh = grids_high.mesh coords_idx = grids_high.coords_idx ngrids0 = numpy.prod(mesh) ngrids1 = grids_high.coords.shape[0] log.debug('mesh %s, ngrids %s/%s', mesh, ngrids1, ngrids0) gx = numpy.fft.fftfreq(mesh[0], 1. / mesh[0]).astype(int) gy = numpy.fft.fftfreq(mesh[1], 1. / mesh[1]).astype(int) gz = numpy.fft.fftfreq(mesh[2], 1. / mesh[2]).astype(int) sub_wvG = wv_freq[:, :, gx[:, None, None], gy[:, None], gz].reshape(-1, ngrids0) wv = tools.ifft(sub_wvG, mesh).real.reshape(2, -1, ngrids0) wv = wv[:, :, coords_idx] if with_j: sub_vG = vG[:, gx[:, None, None], gy[:, None], gz].reshape(-1, ngrids0) vR = tools.ifft(sub_vG, mesh).real.reshape(2, ngrids0) vR = vR[:, coords_idx] idx_h = grids_high.ao_idx if grids_low is None: for ao_h_etc, p0, p1 in mydf.aoR_loop(grids_high, kpts, deriv): ao_h = ao_h_etc[0] add_xc_(veff, ao_h, ao_h, idx_h, idx_h, wv[:, :, p0:p1]) if with_j: add_j_(vj, ao_h, ao_h, idx_h, idx_h, vR[:, p0:p1]) ao_h = ao_h_etc = None else: idx_l = grids_low.ao_idx for ao_h_etc, ao_l_etc in zip( mydf.aoR_loop(grids_high, kpts, deriv), mydf.aoR_loop(grids_low, kpts, deriv)): p0, p1 = ao_h_etc[1:3] ao_h = ao_h_etc[0][0] ao_l = ao_l_etc[0][0] add_xc_(veff, ao_h, ao_h, idx_h, idx_h, wv[:, :, p0:p1]) add_xc_(veff, ao_h, ao_l, idx_h, idx_l, wv[:, :, p0:p1]) add_xc_(veff, ao_l, ao_h, idx_l, idx_h, wv[:, :, p0:p1]) if with_j: add_j_(vj, ao_h, ao_h, idx_h, idx_h, vR[:, p0:p1]) add_j_(vj, ao_h, ao_l, idx_h, idx_l, vR[:, p0:p1]) add_j_(vj, ao_l, ao_h, idx_l, idx_h, vR[:, p0:p1]) ao_h = ao_l = ao_h_etc = ao_l_etc = None vj = _format_jks(vj, dm_kpts, input_band, kpts) veff = _format_jks(veff, dm_kpts, input_band, kpts) return nelec, excsum, veff, vj
def _eval_rhoG(mydf, dm_kpts, hermi=1, kpts=numpy.zeros((1, 3)), deriv=0): log = lib.logger.Logger(mydf.stdout, mydf.verbose) cell = mydf.cell dm_kpts = lib.asarray(dm_kpts, order='C') dms = _format_dms(dm_kpts, kpts) nset, nkpts, nao = dms.shape[:3] tasks = getattr(mydf, 'tasks', None) if tasks is None: mydf.tasks = tasks = multi_grids_tasks(cell, log) log.debug('Multigrid ntasks %s', len(tasks)) assert (deriv <= 2) if abs(dms - dms.transpose(0, 1, 3, 2).conj()).max() < 1e-9: def dot_bra(bra, aodm): rho = numpy.einsum('pi,pi->p', bra.real, aodm.real) if aodm.dtype == numpy.complex: rho += numpy.einsum('pi,pi->p', bra.imag, aodm.imag) return rho if deriv == 0: xctype = 'LDA' rhodim = 1 def make_rho(ao_l, ao_h, dm_lh, dm_hl): c0 = lib.dot(ao_l, dm_lh) rho = dot_bra(ao_h, c0) return rho * 2 elif deriv == 1: xctype = 'GGA' rhodim = 4 def make_rho(ao_l, ao_h, dm_lh, dm_hl): ngrids = ao_l[0].shape[0] rho = numpy.empty((4, ngrids)) c0 = lib.dot(ao_l[0], dm_lh) rho[0] = dot_bra(ao_h[0], c0) for i in range(1, 4): rho[i] = dot_bra(ao_h[i], c0) c0 = lib.dot(ao_h[0], dm_hl) for i in range(1, 4): rho[i] += dot_bra(ao_l[i], c0) return rho * 2 # *2 for dm_lh+dm_hl.T elif deriv == 2: xctype = 'MGGA' rhodim = 6 def make_rho(ao_l, ao_h, dm_lh, dm_hl): ngrids = ao_l[0].shape[0] rho = numpy.empty((6, ngrids)) c = [lib.dot(ao_l[i], dm_lh) for i in range(4)] rho[0] = dot_bra(ao_h[0], c[0]) rho[5] = 0 for i in range(1, 4): rho[i] = dot_bra(ao_h[i], c[0]) rho[i] += dot_bra(ao_h[0], c[i]) rho[5] += dot_bra(ao_h[i], c[i]) * 2 XX, YY, ZZ = 4, 7, 9 ao2 = ao_h[XX] + ao_h[YY] + ao_h[ZZ] rho[4] = dot_bra(ao2, c[0]) ao2 = lib.dot(ao_l[XX] + ao_l[YY] + ao_l[ZZ], dm_lh) rho[4] += dot_bra(ao2, ao_h[0]) rho[4] += rho[5] * 2 rho[5] *= .5 return rho * 2 # *2 for dm_lh+dm_hl.T else: raise NotImplementedError('Non-hermitian density matrices') ni = mydf._numint nx, ny, nz = cell.mesh rhoG = numpy.zeros((nset * rhodim, nx, ny, nz), dtype=numpy.complex) for grids_high, grids_low in tasks: cell_high = grids_high.cell mesh = grids_high.mesh coords_idx = grids_high.coords_idx ngrids0 = numpy.prod(mesh) ngrids1 = grids_high.coords.shape[0] log.debug('mesh %s, ngrids %s/%s', mesh, ngrids1, ngrids0) idx_h = grids_high.ao_idx dms_hh = numpy.asarray(dms[:, :, idx_h[:, None], idx_h], order='C') if grids_low is not None: idx_l = grids_low.ao_idx dms_hl = numpy.asarray(dms[:, :, idx_h[:, None], idx_l], order='C') dms_lh = numpy.asarray(dms[:, :, idx_l[:, None], idx_h], order='C') rho = numpy.zeros((nset, rhodim, ngrids1)) if grids_low is None: for ao_h_etc, p0, p1 in mydf.aoR_loop(grids_high, kpts, deriv): ao_h, mask = ao_h_etc[0], ao_h_etc[2] for k in range(nkpts): for i in range(nset): rho_sub = numint.eval_rho(cell_high, ao_h[k], dms_hh[i, k], mask, xctype, hermi) rho[i, :, p0:p1] += rho_sub.real ao_h = ao_h_etc = None else: for ao_h_etc, ao_l_etc in zip( mydf.aoR_loop(grids_high, kpts, deriv), mydf.aoR_loop(grids_low, kpts, deriv)): p0, p1 = ao_h_etc[1:3] ao_h, mask = ao_h_etc[0][0], ao_h_etc[0][2] ao_l = ao_l_etc[0][0] for k in range(nkpts): for i in range(nset): rho_sub = numint.eval_rho(cell_high, ao_h[k], dms_hh[i, k], mask, xctype, hermi) rho[i, :, p0:p1] += rho_sub.real rho_sub = make_rho(ao_l[k], ao_h[k], dms_lh[i, k], dms_hl[i, k]) rho[i, :, p0:p1] += rho_sub.real ao_h = ao_l = ao_h_etc = ao_l_etc = None rho *= 1. / nkpts rhoR = numpy.zeros((nset * rhodim, ngrids0)) rhoR[:, coords_idx] = rho.reshape(nset * rhodim, ngrids1) gx = numpy.fft.fftfreq(mesh[0], 1. / mesh[0]).astype(int) gy = numpy.fft.fftfreq(mesh[1], 1. / mesh[1]).astype(int) gz = numpy.fft.fftfreq(mesh[2], 1. / mesh[2]).astype(int) rho_freq = tools.fft(rhoR, mesh) * cell.vol / ngrids0 for i in range(nset * rhodim): rhoG[i, gx[:, None, None], gy[:, None], gz] += rho_freq[i].reshape(mesh) return rhoG.reshape(nset, rhodim, ngrids0)
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 mesh = mydf.mesh low_dim_ft_type = mydf.low_dim_ft_type ni = mydf._numint make_rho, nset, nao = ni._gen_rho_evaluator(cell, dm_kpts, hermi) 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, mesh=mesh, low_dim_ft_type=low_dim_ft_type) ngrids = len(coulG) if hermi == 1 or gamma_point(kpts): vR = rhoR = np.zeros((nset, ngrids)) for ao_ks_etc, p0, p1 in mydf.aoR_loop(mydf.grids, kpts): ao_ks, mask = ao_ks_etc[0], ao_ks_etc[2] for i in range(nset): rhoR[i, p0:p1] += make_rho(i, ao_ks, mask, 'LDA') ao = ao_ks = None for i in range(nset): rhoG = tools.fft(rhoR[i], mesh) vG = coulG * rhoG vR[i] = tools.ifft(vG, mesh).real else: # vR may be complex if the underlying density is complex vR = rhoR = np.zeros((nset, ngrids), dtype=np.complex128) for ao_ks_etc, p0, p1 in mydf.aoR_loop(mydf.grids, kpts): ao_ks, mask = ao_ks_etc[0], ao_ks_etc[2] for i in range(nset): for k, ao in enumerate(ao_ks): ao_dm = lib.dot(ao, dms[i, k]) rhoR[i, p0:p1] += np.einsum('xi,xi->x', ao_dm, ao.conj()) rhoR *= 1. / nkpts for i in range(nset): rhoG = tools.fft(rhoR[i], mesh) vG = coulG * rhoG vR[i] = tools.ifft(vG, mesh) kpts_band, input_band = _format_kpts_band(kpts_band, kpts), kpts_band nband = len(kpts_band) weight = cell.vol / ngrids vR *= weight if gamma_point(kpts_band): vj_kpts = np.zeros((nset, nband, nao, nao)) else: vj_kpts = np.zeros((nset, nband, nao, nao), dtype=np.complex128) rho = None for ao_ks_etc, p0, p1 in mydf.aoR_loop(mydf.grids, kpts_band): ao_ks, mask = ao_ks_etc[0], ao_ks_etc[2] for i in range(nset): # ni.eval_mat can handle real vR only # vj_kpts[i] += ni.eval_mat(cell, ao_ks, 1., None, vR[i,p0:p1], mask, 'LDA') for k, ao in enumerate(ao_ks): aow = np.einsum('xi,x->xi', ao, vR[i, p0:p1]) vj_kpts[i, k] += lib.dot(ao.conj().T, aow) return _format_jks(vj_kpts, dm_kpts, input_band, kpts)
cell.gs = numpy.array([n,n,n]) cell.atom = '''C 1.3 .2 .3 C .1 .1 1.1 ''' cell.basis = 'ccpvdz' #cell.basis = {'C': [[0, (2.4, .1, .6), (1.0,.8, .4)], [1, (1.1, 1)]]} #cell.basis = {'C': [[0, (2.4, 1)]]} cell.unit = 'B' #cell.verbose = 4 cell.build(0,0) #cell.nimgs = (2,2,2) ao2 = ft_aopair(cell, cell.Gv) nao = cell.nao_nr() coords = pyscf.pbc.dft.gen_grid.gen_uniform_grids(cell) aoR = pyscf.pbc.dft.numint.eval_ao(cell, coords) aoR2 = numpy.einsum('ki,kj->kij', aoR.conj(), aoR) ngs = aoR.shape[0] for i in range(nao): for j in range(nao): ao2ref = tools.fft(aoR2[:,i,j], cell.gs) * cell.vol/ngs print i, j, numpy.linalg.norm(ao2ref - ao2[:,i,j]) aoG = ft_ao(cell, cell.Gv) for i in range(nao): aoref = tools.fft(aoR[:,i], cell.gs) * cell.vol/ngs print i, numpy.linalg.norm(aoref - aoG[:,i])
cell.mesh = numpy.array([n,n,n]) cell.atom = '''C 1.3 .2 .3 C .1 .1 1.1 ''' cell.basis = 'ccpvdz' #cell.basis = {'C': [[0, (2.4, .1, .6), (1.0,.8, .4)], [1, (1.1, 1)]]} #cell.basis = {'C': [[0, (2.4, 1)]]} cell.unit = 'B' #cell.verbose = 4 cell.build(0,0) #cell.nimgs = (2,2,2) ao2 = ft_aopair(cell, cell.Gv) nao = cell.nao_nr() coords = pyscf.pbc.dft.gen_grid.gen_uniform_grids(cell) aoR = cell.pbc_eval_gto('GTOval', coords) aoR2 = numpy.einsum('ki,kj->kij', aoR.conj(), aoR) ngrids = aoR.shape[0] for i in range(nao): for j in range(nao): ao2ref = tools.fft(aoR2[:,i,j], cell.mesh) * cell.vol/ngrids print(i, j, numpy.linalg.norm(ao2ref - ao2[:,i,j])) aoG = ft_ao(cell, cell.Gv) for i in range(nao): aoref = tools.fft(aoR[:,i], cell.mesh) * cell.vol/ngrids print(i, numpy.linalg.norm(aoref - aoG[:,i]))
def test_ft_aoao(self): #coords = pdft.gen_grid.gen_uniform_grids(cell) #aoR = pdft.numint.eval_ao(cell, coords) #ngrids, nao = aoR.shape #ref = numpy.asarray([tools.fft(aoR[:,i].conj()*aoR[:,j], cell.mesh) # for i in range(nao) for j in range(nao)]) #ref = ref.reshape(nao,nao,-1).transpose(2,0,1) * (cell.vol/ngrids) #dat = ft_ao.ft_aopair(cell, cell.Gv, aosym='s1hermi') #self.assertAlmostEqual(numpy.linalg.norm(ref[:,0,0]-dat[:,0,0]) , 0, 5) #self.assertAlmostEqual(numpy.linalg.norm(ref[:,1,1]-dat[:,1,1]) , 0.02315483195832373, 4) #self.assertAlmostEqual(numpy.linalg.norm(ref[:,2:,2:]-dat[:,2:,2:]), 0, 9) #self.assertAlmostEqual(numpy.linalg.norm(ref[:,0,2:]-dat[:,0,2:]) , 0, 9) #self.assertAlmostEqual(numpy.linalg.norm(ref[:,2:,0]-dat[:,2:,0]) , 0, 9) #idx = numpy.tril_indices(nao) #ref = dat[:,idx[0],idx[1]] #dat = ft_ao.ft_aopair(cell, cell.Gv, aosym='s2') #self.assertAlmostEqual(abs(dat-ref).sum(), 0, 9) coords = pdft.gen_grid.gen_uniform_grids(cell1) Gv, Gvbase, kws = cell1.get_Gv_weights(cell1.mesh) b = cell1.reciprocal_vectors() gxyz = lib.cartesian_prod([numpy.arange(len(x)) for x in Gvbase]) dat = ft_ao.ft_aopair(cell1, cell1.Gv, aosym='s1', b=b, gxyz=gxyz, Gvbase=Gvbase) self.assertAlmostEqual(finger(dat), 1.5666516306798806 + 1.953555017583245j, 9) dat = ft_ao.ft_aopair(cell1, cell1.Gv, aosym='s2', b=b, gxyz=gxyz, Gvbase=Gvbase) self.assertAlmostEqual(finger(dat), -0.85276967757297917 + 1.0378751267506394j, 9) dat = ft_ao.ft_aopair(cell1, cell1.Gv, aosym='s1hermi', b=b, gxyz=gxyz, Gvbase=Gvbase) self.assertAlmostEqual(finger(dat), 1.5666516306798806 + 1.953555017583245j, 9) aoR = pdft.numint.eval_ao(cell1, coords) ngrids, nao = aoR.shape aoaoR = numpy.einsum('pi,pj->ijp', aoR, aoR) ref = tools.fft(aoaoR.reshape(nao * nao, -1), cell1.mesh) ref = ref.reshape(nao, nao, -1).transpose(2, 0, 1) * (cell1.vol / ngrids) self.assertAlmostEqual(numpy.linalg.norm(ref[:, 0, 0] - dat[:, 0, 0]), 0, 7) self.assertAlmostEqual(numpy.linalg.norm(ref[:, 1, 1] - dat[:, 1, 1]), 0, 7) self.assertAlmostEqual( numpy.linalg.norm(ref[:, 2:, 2:] - dat[:, 2:, 2:]), 0, 7) self.assertAlmostEqual( numpy.linalg.norm(ref[:, 0, 2:] - dat[:, 0, 2:]), 0, 7) self.assertAlmostEqual( numpy.linalg.norm(ref[:, 2:, 0] - dat[:, 2:, 0]), 0, 7) idx = numpy.tril_indices(nao) ref = dat[:, idx[0], idx[1]] dat = ft_ao.ft_aopair(cell1, cell1.Gv, aosym='s2') self.assertAlmostEqual(abs(dat - ref).sum(), 0, 9)
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
cell.a = numpy.diag([L, L, L]) cell.gs = numpy.array([n, n, n]) cell.atom = '''C 1.3 .2 .3 C .1 .1 1.1 ''' cell.basis = 'ccpvdz' #cell.basis = {'C': [[0, (2.4, .1, .6), (1.0,.8, .4)], [1, (1.1, 1)]]} #cell.basis = {'C': [[0, (2.4, 1)]]} cell.unit = 'B' #cell.verbose = 4 cell.build(0, 0) #cell.nimgs = (2,2,2) ao2 = ft_aopair(cell, cell.Gv) nao = cell.nao_nr() coords = pyscf.pbc.dft.gen_grid.gen_uniform_grids(cell) aoR = pyscf.pbc.dft.numint.eval_ao(cell, coords) aoR2 = numpy.einsum('ki,kj->kij', aoR.conj(), aoR) ngs = aoR.shape[0] for i in range(nao): for j in range(nao): ao2ref = tools.fft(aoR2[:, i, j], cell.gs) * cell.vol / ngs print(i, j, numpy.linalg.norm(ao2ref - ao2[:, i, j])) aoG = ft_ao(cell, cell.Gv) for i in range(nao): aoref = tools.fft(aoR[:, i], cell.gs) * cell.vol / ngs print(i, numpy.linalg.norm(aoref - aoG[:, i]))
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(aoi, aoj, fac=1): if id(aoi) == id(aoj) and iden_coeffs(mo_coeffs[0], mo_coeffs[1]): moi = moj = numpy.asarray(lib.dot(mo_coeffs[0].T, aoi.T), order='C') else: moi = numpy.asarray(lib.dot(mo_coeffs[0].T, aoi.T), order='C') moj = numpy.asarray(lib.dot(mo_coeffs[1].T, aoj.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 * moi[i].conj() * moj, mydf.gs) mo_pairs_G = mo_pairs_G.reshape(-1, ngs).T return mo_pairs_G if gamma_point(kpts): # gamma point, real ao = mydf._numint.eval_ao(cell, coords, kpts[:1])[0] if compact and iden_coeffs(mo_coeffs[0], mo_coeffs[1]): mo = numpy.asarray(lib.dot(mo_coeffs[0].T, ao.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( mo[i].conj() * mo[:i + 1], mydf.gs) ij += i + 1 mo_pairs_G = mo_pairs_G.T else: mo_pairs_G = trans(ao, ao) elif is_zero(kpts[0] - kpts[1]): ao = mydf._numint.eval_ao(cell, coords, kpts[:1])[0] mo_pairs_G = trans(ao, ao) else: if q is None: q = kpts[1] - kpts[0] aoi, aoj = mydf._numint.eval_ao(cell, coords, kpts) fac = numpy.exp(-1j * numpy.dot(coords, q)) mo_pairs_G = trans(aoi, aoj, fac) return mo_pairs_G
def _make_fftdf_eris(mycc, eris): mydf = mycc._scf.with_df mo_coeff = eris.mo_coeff kpts = mycc.kpts logger = Logger(mycc.stdout, mycc.verbose) cell = mydf.cell gvec = cell.reciprocal_vectors() nao = cell.nao_nr() coords = cell.gen_uniform_grids(mydf.mesh) ngrids = len(coords) nkpts = len(kpts) nocc, nmo = mycc.nocc, mycc.nmo nvir = nmo - nocc cput1 = cput0 = (time.clock(), time.time()) ijG = ctf.zeros([nkpts,nkpts,nocc,nocc,ngrids], dtype=np.complex128) iaG = ctf.zeros([nkpts,nkpts,nocc,nvir,ngrids], dtype=np.complex128) abG = ctf.zeros([nkpts,nkpts,nvir,nvir,ngrids], dtype=np.complex128) ijR = ctf.zeros([nkpts,nkpts,nocc,nocc,ngrids], dtype=np.complex128) iaR = ctf.zeros([nkpts,nkpts,nocc,nvir,ngrids], dtype=np.complex128) aiR = ctf.zeros([nkpts,nkpts,nvir,nocc,ngrids], dtype=np.complex128) abR = ctf.zeros([nkpts,nkpts,nvir,nvir,ngrids], dtype=np.complex128) jobs = [] for ki in range(nkpts): for kj in range(ki,nkpts): jobs.append([ki,kj]) tasks = list(static_partition(jobs)) ntasks = max(comm.allgather(len(tasks))) idx_ooG = np.arange(nocc*nocc*ngrids) idx_ovG = np.arange(nocc*nvir*ngrids) idx_vvG = np.arange(nvir*nvir*ngrids) for itask in range(ntasks): if itask >= len(tasks): ijR.write([], []) iaR.write([], []) aiR.write([], []) abR.write([], []) ijR.write([], []) iaR.write([], []) aiR.write([], []) abR.write([], []) ijG.write([], []) iaG.write([], []) abG.write([], []) ijG.write([], []) iaG.write([], []) abG.write([], []) continue ki, kj = tasks[itask] kpti, kptj = kpts[ki], kpts[kj] ao_kpti = mydf._numint.eval_ao(cell, coords, kpti)[0] ao_kptj = mydf._numint.eval_ao(cell, coords, kptj)[0] q = kptj - kpti coulG = tools.get_coulG(cell, q, mesh=mydf.mesh) wcoulG = coulG * (cell.vol/ngrids) fac = np.exp(-1j * np.dot(coords, q)) mo_kpti = np.dot(ao_kpti, mo_coeff[ki]).T mo_kptj = np.dot(ao_kptj, mo_coeff[kj]).T mo_pairs = np.einsum('ig,jg->ijg', mo_kpti.conj(), mo_kptj) mo_pairs_G = tools.fft(mo_pairs.reshape(-1,ngrids)*fac, mydf.mesh) off = ki * nkpts + kj ijR.write(off*idx_ooG.size+idx_ooG, mo_pairs[:nocc,:nocc].ravel()) iaR.write(off*idx_ovG.size+idx_ovG, mo_pairs[:nocc,nocc:].ravel()) aiR.write(off*idx_ovG.size+idx_ovG, mo_pairs[nocc:,:nocc].ravel()) abR.write(off*idx_vvG.size+idx_vvG, mo_pairs[nocc:,nocc:].ravel()) off = kj * nkpts + ki mo_pairs = mo_pairs.transpose(1,0,2).conj() ijR.write(off*idx_ooG.size+idx_ooG, mo_pairs[:nocc,:nocc].ravel()) iaR.write(off*idx_ovG.size+idx_ovG, mo_pairs[:nocc,nocc:].ravel()) aiR.write(off*idx_ovG.size+idx_ovG, mo_pairs[nocc:,:nocc].ravel()) abR.write(off*idx_vvG.size+idx_vvG, mo_pairs[nocc:,nocc:].ravel()) mo_pairs = None mo_pairs_G*= wcoulG v = tools.ifft(mo_pairs_G, mydf.mesh) v *= fac.conj() v = v.reshape(nmo,nmo,ngrids) off = ki * nkpts + kj ijG.write(off*idx_ooG.size+idx_ooG, v[:nocc,:nocc].ravel()) iaG.write(off*idx_ovG.size+idx_ovG, v[:nocc,nocc:].ravel()) abG.write(off*idx_vvG.size+idx_vvG, v[nocc:,nocc:].ravel()) off = kj * nkpts + ki v = v.transpose(1,0,2).conj() ijG.write(off*idx_ooG.size+idx_ooG, v[:nocc,:nocc].ravel()) iaG.write(off*idx_ovG.size+idx_ovG, v[:nocc,nocc:].ravel()) abG.write(off*idx_vvG.size+idx_vvG, v[nocc:,nocc:].ravel()) cput1 = logger.timer("Generating ijG", *cput1) sym1 = ["+-+", [kpts,]*3, None, gvec] sym2 = ["+--", [kpts,]*3, None, gvec] ooG = tensor(ijG, sym1, verbose=mycc.SYMVERBOSE) ovG = tensor(iaG, sym1, verbose=mycc.SYMVERBOSE) vvG = tensor(abG, sym1, verbose=mycc.SYMVERBOSE) ooR = tensor(ijR, sym2, verbose=mycc.SYMVERBOSE) ovR = tensor(iaR, sym2, verbose=mycc.SYMVERBOSE) voR = tensor(aiR, sym2, verbose=mycc.SYMVERBOSE) vvR = tensor(abR, sym2, verbose=mycc.SYMVERBOSE) eris.oooo = lib.einsum('ijg,klg->ijkl', ooG, ooR)/ nkpts eris.ooov = lib.einsum('ijg,kag->ijka', ooG, ovR)/ nkpts eris.oovv = lib.einsum('ijg,abg->ijab', ooG, vvR)/ nkpts ooG = ooR = ijG = ijR = None eris.ovvo = lib.einsum('iag,bjg->iabj', ovG, voR)/ nkpts eris.ovov = lib.einsum('iag,jbg->iajb', ovG, ovR)/ nkpts ovR = iaR = voR = aiR = None eris.ovvv = lib.einsum('iag,bcg->iabc', ovG, vvR)/ nkpts ovG = iaG = None eris.vvvv = lib.einsum('abg,cdg->abcd', vvG, vvR)/ nkpts cput1 = logger.timer("ijG to eri", *cput1)