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 get_nuc(mydf, kpts): mydf = _sync_mydf(mydf) cell = mydf.cell if kpts is None: kpts_lst = numpy.zeros((1,3)) else: kpts_lst = numpy.reshape(kpts, (-1,3)) if abs(kpts_lst).sum() < 1e-9: # gamma_point dtype = numpy.float64 else: dtype = numpy.complex128 gs = mydf.gs charge = -cell.atom_charges() Gv = cell.get_Gv(gs) SI = cell.get_SI(Gv) rhoG = numpy.dot(charge, SI) coulG = tools.get_coulG(cell, gs=gs, Gv=Gv) vneG = rhoG * coulG vneR = tools.ifft(vneG, mydf.gs).real vne = [lib.dot(aoR.T.conj()*vneR, aoR) for k, aoR in mydf.mpi_aoR_loop(gs, kpts_lst)] vne = mpi.gather(lib.asarray(vne, dtype=dtype)) if rank == 0: if kpts is None or numpy.shape(kpts) == (3,): vne = vne[0] return vne
def test_pnucp(self): cell1 = gto.Cell() cell1.atom = ''' He 1.3 .2 .3 He .1 .1 1.1 ''' cell1.basis = {'He': [[0, [0.8, 1]], [1, [0.6, 1]] ]} cell1.mesh = [15]*3 cell1.a = numpy.array(([2.0, .9, 0. ], [0.1, 1.9, 0.4], [0.8, 0 , 2.1])) cell1.build() charge = -cell1.atom_charges() Gv = cell1.get_Gv(cell1.mesh) SI = cell1.get_SI(Gv) rhoG = numpy.dot(charge, SI) coulG = tools.get_coulG(cell1, mesh=cell1.mesh, Gv=Gv) vneG = rhoG * coulG vneR = tools.ifft(vneG, cell1.mesh).real coords = cell1.gen_uniform_grids(cell1.mesh) aoR = dft.numint.eval_ao(cell1, coords, deriv=1) ngrids, nao = aoR.shape[1:] vne_ref = numpy.einsum('p,xpi,xpj->ij', vneR, aoR[1:4], aoR[1:4]) mydf = df.AFTDF(cell1) dat = sfx2c1e.get_pnucp(mydf) self.assertAlmostEqual(abs(dat-vne_ref).max(), 0, 7) mydf.eta = 0 dat = sfx2c1e.get_pnucp(mydf) self.assertAlmostEqual(abs(dat-vne_ref).max(), 0, 7)
def get_nuc(mydf, kpts): mydf = _sync_mydf(mydf) cell = mydf.cell if kpts is None: kpts_lst = numpy.zeros((1, 3)) else: kpts_lst = numpy.reshape(kpts, (-1, 3)) if abs(kpts_lst).sum() < 1e-9: # gamma_point dtype = numpy.float64 else: dtype = numpy.complex128 mesh = mydf.mesh charge = -cell.atom_charges() Gv = cell.get_Gv(mesh) SI = cell.get_SI(Gv) rhoG = numpy.dot(charge, SI) coulG = tools.get_coulG(cell, mesh=mesh, Gv=Gv) vneG = rhoG * coulG vneR = tools.ifft(vneG, mydf.mesh).real vne = [0] * len(kpts_lst) for ao_ks_etc, p0, p1 in mydf.mpi_aoR_loop(mydf.grids, kpts_lst): ao_ks = ao_ks_etc[0] for k, ao in enumerate(ao_ks): vne[k] += lib.dot(ao.T.conj() * vneR[p0:p1], ao) ao = ao_ks = None vne = mpi.reduce(lib.asarray(vne)) if rank == 0: if kpts is None or numpy.shape(kpts) == (3, ): vne = vne[0] return vne
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 test_pnucp(self): cell1 = gto.Cell() cell1.atom = ''' He 1.3 .2 .3 He .1 .1 1.1 ''' cell1.basis = {'He': [[0, [0.8, 1]], [1, [0.6, 1]]]} cell1.mesh = [15] * 3 cell1.a = numpy.array(([2.0, .9, 0.], [0.1, 1.9, 0.4], [0.8, 0, 2.1])) cell1.build() charge = -cell1.atom_charges() Gv = cell1.get_Gv(cell1.mesh) SI = cell1.get_SI(Gv) rhoG = numpy.dot(charge, SI) coulG = tools.get_coulG(cell1, mesh=cell1.mesh, Gv=Gv) vneG = rhoG * coulG vneR = tools.ifft(vneG, cell1.mesh).real coords = cell1.gen_uniform_grids(cell1.mesh) aoR = dft.numint.eval_ao(cell1, coords, deriv=1) ngrids, nao = aoR.shape[1:] vne_ref = numpy.einsum('p,xpi,xpj->ij', vneR, aoR[1:4], aoR[1:4]) mydf = df.AFTDF(cell1) dat = sfx2c1e.get_pnucp(mydf) self.assertAlmostEqual(abs(dat - vne_ref).max(), 0, 7) mydf.eta = 0 dat = sfx2c1e.get_pnucp(mydf) self.assertAlmostEqual(abs(dat - vne_ref).max(), 0, 7)
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 get_nuc(mydf, kpts=None): if kpts is None: kpts_lst = numpy.zeros((1, 3)) else: kpts_lst = numpy.reshape(kpts, (-1, 3)) cell = mydf.cell mesh = mydf.mesh charge = -cell.atom_charges() Gv = cell.get_Gv(mesh) SI = cell.get_SI(Gv) rhoG = numpy.dot(charge, SI) coulG = tools.get_coulG(cell, mesh=mesh, Gv=Gv) vneG = rhoG * coulG vneR = tools.ifft(vneG, mesh).real vne = [0] * len(kpts_lst) for ao_ks_etc, p0, p1 in mydf.aoR_loop(mydf.grids, kpts_lst): ao_ks = ao_ks_etc[0] for k, ao in enumerate(ao_ks): vne[k] += lib.dot(ao.T.conj() * vneR[p0:p1], ao) ao = ao_ks = None if kpts is None or numpy.shape(kpts) == (3, ): vne = vne[0] return numpy.asarray(vne)
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)) gs = mydf.gs charge = -cell.atom_charges() Gv = cell.get_Gv(gs) SI = cell.get_SI(Gv) rhoG = numpy.dot(charge, SI) coulG = tools.get_coulG(cell, gs=gs, Gv=Gv) vneG = rhoG * coulG vneR = tools.ifft(vneG, mydf.gs).real vne = [ lib.dot(aoR.T.conj() * vneR, aoR) for k, aoR in mydf.aoR_loop(gs, kpts_lst) ] if kpts is None or numpy.shape(kpts) == (3, ): vne = vne[0] return numpy.asarray(vne)
def test_ifft(self): n = 31 a = numpy.random.random([2,n,n,n]) ref = numpy.fft.ifftn(a, axes=(1,2,3)).ravel() v = tools.ifft(a, [n,n,n]).ravel() self.assertAlmostEqual(abs(ref-v).max(), 0, 10) a = numpy.random.random([2,n,n,8]) ref = numpy.fft.ifftn(a, axes=(1,2,3)).ravel() v = tools.ifft(a, [n,n,8]).ravel() self.assertAlmostEqual(abs(ref-v).max(), 0, 10) a = numpy.random.random([2,8,n,8]) ref = numpy.fft.ifftn(a, axes=(1,2,3)).ravel() v = tools.ifft(a, [8,n,8]).ravel() self.assertAlmostEqual(abs(ref-v).max(), 0, 10)
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) 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 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): vj_kpts[i] += ni.eval_mat(cell, ao_ks, 1., None, vR[i, p0:p1], mask, 'LDA') return _format_jks(vj_kpts, dm_kpts, input_band, kpts)
def get_hcore(cell, kpts): '''Part of the nuclear gradients of core Hamiltonian''' h1 = np.asarray(cell.pbc_intor('int1e_ipkin', kpts=kpts)) dtype = h1.dtype if cell._pseudo: SI = cell.get_SI() nao = cell.nao_nr() Gv = cell.Gv natom = cell.natm coords = cell.get_uniform_grids() ngrids, nkpts = len(coords), len(kpts) vlocG = get_vlocG(cell) vpplocG = -np.einsum('ij,ij->j', SI, vlocG) vpplocG[0] = np.sum(get_alphas(cell)) vpplocR = tools.ifft(vpplocG, cell.mesh).real fakemol = _make_fakemol() ptr = mole.PTR_ENV_START for kn, kpt in enumerate(kpts): aos = eval_ao_kpts(cell, coords, kpt, deriv=1)[0] vloc = np.einsum('agi,g,gj->aij', aos[1:].conj(), vpplocR, aos[0]) expir = np.exp(-1j * np.dot(coords, kpt)) aokG = np.asarray([ tools.fftk(np.asarray(ao.T, order='C'), cell.mesh, expir).T for ao in aos ]) Gk = Gv + kpt G_rad = lib.norm(Gk, axis=1) vnl = np.zeros(vloc.shape, dtype=np.complex128) for ia in range(natom): symb = cell.atom_symbol(ia) if symb not in cell._pseudo: continue pp = cell._pseudo[symb] for l, proj in enumerate(pp[5:]): rl, nl, hl = proj if nl > 0: hl = np.asarray(hl) fakemol._bas[0, mole.ANG_OF] = l fakemol._env[ptr + 3] = .5 * rl**2 fakemol._env[ptr + 4] = rl**(l + 1.5) * np.pi**1.25 pYlm_part = fakemol.eval_gto('GTOval', Gk) pYlm = np.empty((nl, l * 2 + 1, ngrids)) for k in range(nl): qkl = _qli(G_rad * rl, l, k) pYlm[k] = pYlm_part.T * qkl SPG_lmi = np.einsum('g,nmg->nmg', SI[ia].conj(), pYlm) SPG_lm_aoG = np.einsum('nmg,agp->anmp', SPG_lmi, aokG) tmp = np.einsum('ij,ajmp->aimp', hl, SPG_lm_aoG[1:]) vnl += np.einsum('aimp,imq->apq', tmp.conj(), SPG_lm_aoG[0]) vnl *= (1. / ngrids**2) if dtype == np.float64: h1[kn, :] += vloc.real + vnl.real else: h1[kn, :] += vloc + vnl else: raise NotImplementedError return h1
def get_pp_o1(cell, kpt=np.zeros(3)): '''Get the periodic pseudotential nuc-el AO matrix, with G=0 removed. ''' coords = pyscf.pbc.dft.gen_grid.gen_uniform_grids(cell) aoR = pyscf.pbc.dft.numint.eval_ao(cell, coords, kpt) nao = cell.nao_nr() SI = cell.get_SI() vlocG = pseudo.get_vlocG(cell) vpplocG = -np.sum(SI * vlocG, axis=0) # vpploc evaluated in real-space vpplocR = tools.ifft(vpplocG, cell.gs) vpploc = np.dot(aoR.T.conj(), vpplocR.reshape(-1,1)*aoR) # vppnonloc evaluated in reciprocal space aokG = np.empty(aoR.shape, np.complex128) for i in range(nao): aokG[:,i] = tools.fftk(aoR[:,i], cell.gs, coords, kpt) ngs = len(aokG) fakemol = pyscf.gto.Mole() fakemol._atm = np.zeros((1,pyscf.gto.ATM_SLOTS), dtype=np.int32) fakemol._bas = np.zeros((1,pyscf.gto.BAS_SLOTS), dtype=np.int32) fakemol._env = np.zeros(5) fakemol._bas[0,pyscf.gto.NPRIM_OF ] = 1 fakemol._bas[0,pyscf.gto.NCTR_OF ] = 1 fakemol._bas[0,pyscf.gto.PTR_EXP ] = 3 fakemol._bas[0,pyscf.gto.PTR_COEFF] = 4 Gv = np.asarray(cell.Gv+kpt) G_rad = pyscf.lib.norm(Gv, axis=1) vppnl = np.zeros((nao,nao), dtype=np.complex128) for ia in range(cell.natm): pp = cell._pseudo[cell.atom_symbol(ia)] for l, proj in enumerate(pp[5:]): rl, nl, hl = proj if nl > 0: hl = np.asarray(hl) fakemol._bas[0,pyscf.gto.ANG_OF] = l fakemol._env[3] = .5*rl**2 fakemol._env[4] = rl**(l+1.5)*np.pi**1.25 pYlm_part = pyscf.dft.numint.eval_ao(fakemol, Gv, deriv=0) pYlm = np.empty((nl,l*2+1,ngs)) for k in range(nl): qkl = pseudo.pp._qli(G_rad*rl, l, k) pYlm[k] = pYlm_part.T * qkl # pYlm is real SPG_lmi = np.einsum('g,nmg->nmg', SI[ia].conj(), pYlm) SPG_lm_aoG = np.einsum('nmg,gp->nmp', SPG_lmi, aokG) tmp = np.einsum('ij,jmp->imp', hl, SPG_lm_aoG) vppnl += np.einsum('imp,imq->pq', SPG_lm_aoG.conj(), tmp) vppnl *= (1./ngs**2) return vpploc + vppnl
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./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 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 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 test_eval_rhoG_nonorth_gga(self): mydf = multigrid.MultiGridFFTDF(cell_nonorth) rhoG = multigrid._eval_rhoG(mydf, dm, hermi=1, kpts=kpts, deriv=1, rhog_high_order=True) mydf = df.FFTDF(cell_nonorth) ni = dft.numint.KNumInt() ao_kpts = ni.eval_ao(cell_nonorth, mydf.grids.coords, kpts, deriv=1) ref = ni.eval_rho(cell_nonorth, ao_kpts, dm, xctype='GGA') rhoR = tools.ifft(rhoG[0], cell_nonorth.mesh).real rhoR *= numpy.prod(cell_nonorth.mesh)/cell_nonorth.vol self.assertAlmostEqual(abs(rhoR-ref).max(), 0, 5)
def hcore_deriv(atm_id): shl0, shl1, p0, p1 = aoslices[atm_id] symb = cell.atom_symbol(atm_id) fakemol = _make_fakemol() vloc_g = 1j * np.einsum('ga,g->ag', Gv, SI[atm_id] * vlocG[atm_id]) nkpts, nao = h1.shape[0], h1.shape[2] hcore = np.zeros([3, nkpts, nao, nao], dtype=h1.dtype) for kn, kpt in enumerate(kpts): ao = eval_ao_kpts(cell, coords, kpt)[0] rho = np.einsum('gi,gj->gij', ao.conj(), ao) for ax in range(3): vloc_R = tools.ifft(vloc_g[ax], mesh).real vloc = np.einsum('gij,g->ij', rho, vloc_R) hcore[ax, kn] += vloc rho = None aokG = tools.fftk(np.asarray(ao.T, order='C'), mesh, np.exp(-1j * np.dot(coords, kpt))).T ao = None Gk = Gv + kpt G_rad = lib.norm(Gk, axis=1) if symb not in cell._pseudo: continue pp = cell._pseudo[symb] for l, proj in enumerate(pp[5:]): rl, nl, hl = proj if nl > 0: hl = np.asarray(hl) fakemol._bas[0, mole.ANG_OF] = l fakemol._env[ptr + 3] = .5 * rl**2 fakemol._env[ptr + 4] = rl**(l + 1.5) * np.pi**1.25 pYlm_part = fakemol.eval_gto('GTOval', Gk) pYlm = np.empty((nl, l * 2 + 1, ngrids)) for k in range(nl): qkl = _qli(G_rad * rl, l, k) pYlm[k] = pYlm_part.T * qkl SPG_lmi = np.einsum('g,nmg->nmg', SI[atm_id].conj(), pYlm) SPG_lm_aoG = np.einsum('nmg,gp->nmp', SPG_lmi, aokG) SPG_lmi_G = 1j * np.einsum('nmg, ga->anmg', SPG_lmi, Gv) SPG_lm_G_aoG = np.einsum('anmg, gp->anmp', SPG_lmi_G, aokG) tmp_1 = np.einsum('ij,ajmp->aimp', hl, SPG_lm_G_aoG) tmp_2 = np.einsum('ij,jmp->imp', hl, SPG_lm_aoG) vppnl = np.einsum( 'imp,aimq->apq', SPG_lm_aoG.conj(), tmp_1) + np.einsum( 'aimp,imq->apq', SPG_lm_G_aoG.conj(), tmp_2) vppnl *= (1. / ngrids**2) if dtype == np.float64: hcore[:, kn] += vppnl.real else: hcore[:, kn] += vppnl hcore[:, kn, p0:p1] -= h1[kn, :, p0:p1] hcore[:, kn, :, p0:p1] -= h1[kn, :, p0:p1].transpose(0, 2, 1).conj() return hcore
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 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, single_kpt_band = _format_kpts_band(kpts_band, kpts) nband = len(kpts_band) vj_kpts = [] 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, kpts_band, kpts, single_kpt_band)
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_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 test_eval_rhoG_orth_kpts(self): numpy.random.seed(9) dm = numpy.random.random(dm1.shape) + numpy.random.random(dm1.shape) * 1j mydf = multigrid.MultiGridFFTDF(cell_orth) rhoG = multigrid._eval_rhoG(mydf, dm, hermi=0, kpts=kpts, deriv=0, rhog_high_order=True) self.assertTrue(rhoG.dtype == numpy.complex128) mydf = df.FFTDF(cell_orth) ni = dft.numint.KNumInt() ao_kpts = ni.eval_ao(cell_orth, mydf.grids.coords, kpts, deriv=0) ref = ni.eval_rho(cell_orth, ao_kpts, dm, hermi=0, xctype='LDA') rhoR = tools.ifft(rhoG[0], cell_orth.mesh).real rhoR *= numpy.prod(cell_orth.mesh)/cell_orth.vol self.assertAlmostEqual(abs(rhoR-ref).max(), 0, 7)
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 get_nuc(cell, kpt=np.zeros(3)): '''Get the bare periodic nuc-el AO matrix, with G=0 removed. See Martin (12.16)-(12.21). ''' coords = gen_grid.gen_uniform_grids(cell) aoR = numint.eval_ao(cell, coords, kpt) chargs = cell.atom_charges() SI = cell.get_SI() coulG = tools.get_coulG(cell) vneG = -np.dot(chargs, SI) * coulG vneR = tools.ifft(vneG, cell.gs).real vne = np.dot(aoR.T.conj(), vneR.reshape(-1, 1) * aoR) return vne
def test_eval_rhoG_nonorth_gga(self): mydf = multigrid.MultiGridFFTDF(cell_nonorth) rhoG = multigrid._eval_rhoG(mydf, dm, hermi=1, kpts=kpts, deriv=1, rhog_high_order=True) mydf = df.FFTDF(cell_nonorth) ni = dft.numint.KNumInt() ao_kpts = ni.eval_ao(cell_nonorth, mydf.grids.coords, kpts, deriv=1) ref = ni.eval_rho(cell_nonorth, ao_kpts, dm, xctype='GGA') rhoR = tools.ifft(rhoG[0], cell_nonorth.mesh).real rhoR *= numpy.prod(cell_nonorth.mesh) / cell_nonorth.vol self.assertAlmostEqual(abs(rhoR - ref).max(), 0, 7)
def get_nuc(cell, kpt=np.zeros(3)): '''Get the bare periodic nuc-el AO matrix, with G=0 removed. See Martin (12.16)-(12.21). ''' coords = pyscf.pbc.dft.gen_grid.gen_uniform_grids(cell) aoR = pyscf.pbc.dft.numint.eval_ao(cell, coords, kpt) chargs = [cell.atom_charge(i) for i in range(cell.natm)] SI = cell.get_SI() coulG = tools.get_coulG(cell) vneG = -np.dot(chargs,SI) * coulG vneR = tools.ifft(vneG, cell.gs) vne = np.dot(aoR.T.conj(), vneR.reshape(-1,1)*aoR) return vne
def get_j_kpts(mydf, dm_kpts, hermi=1, kpts=numpy.zeros((1,3)), kpts_band=None): mydf = _sync_mydf(mydf) cell = mydf.cell mesh = mydf.mesh 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) vR = rhoR = numpy.zeros((nset,ngrids)) for ao_ks_etc, p0, p1 in mydf.mpi_aoR_loop(mydf.grids, kpts): ao_ks = ao_ks_etc[0] for k, ao in enumerate(ao_ks): for i in range(nset): rhoR[i,p0:p1] += numint.eval_rho(cell, ao, dms[i,k]) ao = ao_ks = None rhoR = mpi.allreduce(rhoR) for i in range(nset): rhoR[i] *= 1./nkpts rhoG = tools.fft(rhoR[i], mesh) vG = coulG * rhoG vR[i] = tools.ifft(vG, mesh).real 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 = numpy.zeros((nset,nband,nao,nao)) else: vj_kpts = numpy.zeros((nset,nband,nao,nao), dtype=numpy.complex128) for ao_ks_etc, p0, p1 in mydf.mpi_aoR_loop(mydf.grids, kpts_band): ao_ks = ao_ks_etc[0] for k, ao in enumerate(ao_ks): for i in range(nset): vj_kpts[i,k] += lib.dot(ao.T.conj()*vR[i,p0:p1], ao) vj_kpts = mpi.reduce(vj_kpts) if gamma_point(kpts_band): vj_kpts = vj_kpts.real return _format_jks(vj_kpts, dm_kpts, input_band, kpts)
def get_pp(cell, kpt=np.zeros(3)): '''Get the periodic pseudotential nuc-el AO matrix, with G=0 removed. ''' coords = gen_grid.gen_uniform_grids(cell) aoR = numint.eval_ao(cell, coords, kpt) nao = cell.nao_nr() SI = cell.get_SI() vlocG = pseudo.get_vlocG(cell) vlocG[:, 0] = 0 vpplocG = -np.sum(SI * vlocG, axis=0) # vpploc evaluated in real-space vpplocR = tools.ifft(vpplocG, cell.mesh) vpploc = np.dot(aoR.T.conj(), vpplocR.reshape(-1, 1) * aoR) # vppnonloc evaluated in reciprocal space aokG = np.empty(aoR.shape, np.complex128) expmikr = np.exp(-1j * np.dot(coords, kpt)) for i in range(nao): aokG[:, i] = tools.fftk(aoR[:, i], cell.mesh, expmikr) ngrids = len(aokG) vppnl = np.zeros((nao, nao), dtype=np.complex128) hs, projGs = pseudo.get_projG(cell, kpt) for ia, [h_ia, projG_ia] in enumerate(zip(hs, projGs)): for l, h in enumerate(h_ia): nl = h.shape[0] for m in range(-l, l + 1): SPG_lm_aoG = np.zeros((nl, nao), dtype=np.complex128) for i in range(nl): SPG_lmi = SI[ia, :] * projG_ia[l][m][i] SPG_lm_aoG[i, :] = np.einsum('g,gp->p', SPG_lmi.conj(), aokG) for i in range(nl): for j in range(nl): # Note: There is no (-1)^l here. vppnl += h[i, j] * np.einsum('p,q->pq', SPG_lm_aoG[i, :].conj(), SPG_lm_aoG[j, :]) vppnl *= (1. / ngrids**2) ovlp = cell.pbc_intor('int1e_ovlp_sph', hermi=1, kpts=kpt) vpploc += 1. / cell.vol * np.sum(pseudo.get_alphas(cell)) * ovlp return vpploc + vppnl
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 get_pp(cell, kpt=np.zeros(3)): '''Get the periodic pseudotential nuc-el AO matrix, with G=0 removed. ''' coords = gen_grid.gen_uniform_grids(cell) aoR = numint.eval_ao(cell, coords, kpt) nao = cell.nao_nr() SI = cell.get_SI() vlocG = pseudo.get_vlocG(cell) vlocG[:,0] = 0 vpplocG = -np.sum(SI * vlocG, axis=0) # vpploc evaluated in real-space vpplocR = tools.ifft(vpplocG, cell.mesh) vpploc = np.dot(aoR.T.conj(), vpplocR.reshape(-1,1)*aoR) # vppnonloc evaluated in reciprocal space aokG = np.empty(aoR.shape, np.complex128) expmikr = np.exp(-1j*np.dot(coords,kpt)) for i in range(nao): aokG[:,i] = tools.fftk(aoR[:,i], cell.mesh, expmikr) ngrids = len(aokG) vppnl = np.zeros((nao,nao), dtype=np.complex128) hs, projGs = pseudo.get_projG(cell, kpt) for ia, [h_ia,projG_ia] in enumerate(zip(hs,projGs)): for l, h in enumerate(h_ia): nl = h.shape[0] for m in range(-l,l+1): SPG_lm_aoG = np.zeros((nl,nao), dtype=np.complex128) for i in range(nl): SPG_lmi = SI[ia,:] * projG_ia[l][m][i] SPG_lm_aoG[i,:] = np.einsum('g,gp->p', SPG_lmi.conj(), aokG) for i in range(nl): for j in range(nl): # Note: There is no (-1)^l here. vppnl += h[i,j]*np.einsum('p,q->pq', SPG_lm_aoG[i,:].conj(), SPG_lm_aoG[j,:]) vppnl *= (1./ngrids**2) ovlp = cell.pbc_intor('int1e_ovlp_sph', hermi=1, kpts=kpt) vpploc += 1./cell.vol * np.sum(pseudo.get_alphas(cell)) * ovlp return vpploc + vppnl
def get_jmod_pw_poisson(cell, modcell, kpt=None): modchg = np.asarray([cell.atom_charge(ia) for ia in range(cell.natm)]) rhok = 0 k2 = np.einsum('ij,ij->i', cell.Gv, cell.Gv) for ib in range(modcell.nbas): e = modcell.bas_exp(ib)[0] r = modcell.bas_coord(ib) si = np.exp(-1j*np.einsum('ij,j->i', cell.Gv, r)) rhok += modchg[ib] * si * np.exp(-k2/(4*e)) vk = rhok * tools.get_coulG(cell) # weight = vol/N, 1/vol * weight = 1/N # ifft has 1/N vw = tools.ifft(vk, cell.gs) coords = pdft.gen_grid.gen_uniform_grids(cell) aoR = pdft.numint.eval_ao(cell, coords, None) vj = np.dot(aoR.T.conj(), vw.reshape(-1,1) * aoR) return vj
def test_eval_rhoG_orth_kpts(self): numpy.random.seed(9) dm = numpy.random.random( dm1.shape) + numpy.random.random(dm1.shape) * 1j mydf = multigrid.MultiGridFFTDF(cell_orth) rhoG = multigrid._eval_rhoG(mydf, dm, hermi=0, kpts=kpts, deriv=0, rhog_high_order=True) self.assertTrue(rhoG.dtype == numpy.complex128) mydf = df.FFTDF(cell_orth) ni = dft.numint.KNumInt() ao_kpts = ni.eval_ao(cell_orth, mydf.grids.coords, kpts, deriv=0) ref = ni.eval_rho(cell_orth, ao_kpts, dm, hermi=0, xctype='LDA') rhoR = tools.ifft(rhoG[0], cell_orth.mesh).real rhoR *= numpy.prod(cell_orth.mesh) / cell_orth.vol self.assertAlmostEqual(abs(rhoR - ref).max(), 0, 8)
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 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)) gs = mydf.gs charge = -cell.atom_charges() Gv = cell.get_Gv(gs) SI = cell.get_SI(Gv) rhoG = numpy.dot(charge, SI) coulG = tools.get_coulG(cell, gs=gs, Gv=Gv) vneG = rhoG * coulG vneR = tools.ifft(vneG, mydf.gs).real vne = [lib.dot(aoR.T.conj() * vneR, aoR) for k, aoR in mydf.aoR_loop(gs, kpts_lst)] if kpts is None or numpy.shape(kpts) == (3,): vne = vne[0] return vne
def get_pp_loc_part2(cell, kpt=np.zeros(3)): coords = gen_grid.gen_uniform_grids(cell) aoR = numint.eval_ao(cell, coords, kpt) nao = cell.nao_nr() SI = cell.get_SI() G = lib.norm(cell.Gv, axis=1) vlocG = np.zeros((cell.natm,len(G))) for ia in range(cell.natm): Zia = cell.atom_charge(ia) symb = cell.atom_symbol(ia) if symb not in cell._pseudo: vlocG[ia] = 0 continue pp = cell._pseudo[symb] rloc, nexp, cexp = pp[1:3+1] G_red = G*rloc cfacs = np.array( [1*G_red**0, 3 - G_red**2, 15 - 10*G_red**2 + G_red**4, 105 - 105*G_red**2 + 21*G_red**4 - G_red**6]) with np.errstate(divide='ignore'): # Note the signs -- potential here is positive vlocG[ia,:] = (# 4*np.pi * Zia * np.exp(-0.5*G_red**2)/G**2 - (2*np.pi)**(3/2.)*rloc**3*np.exp(-0.5*G_red**2)*( np.dot(cexp, cfacs[:nexp])) ) vpplocG = -np.sum(SI * vlocG, axis=0) vpplocR = tools.ifft(vpplocG, cell.gs).real vpploc = np.dot(aoR.T.conj(), vpplocR.reshape(-1,1)*aoR) if aoR.dtype == np.double: return vpploc.real else: return vpploc
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 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) nao = cell.nao_nr() # 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 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] for l, proj in enumerate(pp[5:]): rl, nl, hl = proj if nl > 0: hl = numpy.asarray(hl) fakemol._bas[0, gto.ANG_OF] = l fakemol._env[ptr + 3] = 0.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) pYlm = numpy.empty((nl, l * 2 + 1, ngs)) for k in range(nl): qkl = pseudo.pp._qli(G_rad * rl, l, k) pYlm[k] = pYlm_part.T * qkl # pYlm is real 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) return vppnl * (1.0 / ngs ** 2) for k, kpt in enumerate(kpts_lst): vppnl = vppnl_by_k(kpt) if abs(kpt).sum() < 1e-9: # gamma_point vpp[k] = vpp[k].real + vppnl.real else: vpp[k] += vppnl if kpts is None or numpy.shape(kpts) == (3,): vpp = vpp[0] return vpp
def get_pp(cell, kpt=np.zeros(3)): '''Get the periodic pseudotential nuc-el AO matrix ''' import pyscf.dft from pyscf.pbc import tools from pyscf.pbc.dft import gen_grid from pyscf.pbc.dft import numint coords = gen_grid.gen_uniform_grids(cell) aoR = numint.eval_ao(cell, coords, kpt) nao = cell.nao_nr() SI = cell.get_SI() vlocG = get_vlocG(cell) vpplocG = -np.sum(SI * vlocG, axis=0) vpplocG[0] = np.sum(get_alphas(cell)) # from get_jvloc_G0 function # vpploc evaluated in real-space vpplocR = tools.ifft(vpplocG, cell.gs).real vpploc = np.dot(aoR.T.conj(), vpplocR.reshape(-1, 1) * aoR) # vppnonloc evaluated in reciprocal space aokG = tools.fftk(np.asarray(aoR.T, order='C'), cell.gs, np.exp(-1j * np.dot(coords, kpt))).T ngs = len(aokG) fakemol = pyscf.gto.Mole() fakemol._atm = np.zeros((1, pyscf.gto.ATM_SLOTS), dtype=np.int32) fakemol._bas = np.zeros((1, pyscf.gto.BAS_SLOTS), dtype=np.int32) ptr = pyscf.gto.PTR_ENV_START fakemol._env = np.zeros(ptr + 10) fakemol._bas[0, pyscf.gto.NPRIM_OF] = 1 fakemol._bas[0, pyscf.gto.NCTR_OF] = 1 fakemol._bas[0, pyscf.gto.PTR_EXP] = ptr + 3 fakemol._bas[0, pyscf.gto.PTR_COEFF] = ptr + 4 Gv = np.asarray(cell.Gv + kpt) G_rad = lib.norm(Gv, axis=1) vppnl = np.zeros((nao, nao), dtype=np.complex128) for ia in range(cell.natm): symb = cell.atom_symbol(ia) if symb not in cell._pseudo: continue pp = cell._pseudo[symb] for l, proj in enumerate(pp[5:]): rl, nl, hl = proj if nl > 0: hl = np.asarray(hl) fakemol._bas[0, pyscf.gto.ANG_OF] = l fakemol._env[ptr + 3] = .5 * rl**2 fakemol._env[ptr + 4] = rl**(l + 1.5) * np.pi**1.25 pYlm_part = pyscf.dft.numint.eval_ao(fakemol, Gv, deriv=0) pYlm = np.empty((nl, l * 2 + 1, ngs)) for k in range(nl): qkl = _qli(G_rad * rl, l, k) pYlm[k] = pYlm_part.T * qkl # pYlm is real SPG_lmi = np.einsum('g,nmg->nmg', SI[ia].conj(), pYlm) SPG_lm_aoG = np.einsum('nmg,gp->nmp', SPG_lmi, aokG) tmp = np.einsum('ij,jmp->imp', hl, SPG_lm_aoG) vppnl += np.einsum('imp,imq->pq', SPG_lm_aoG.conj(), tmp) vppnl *= (1. / ngs**2) if aoR.dtype == np.double: return vpploc.real + vppnl.real else: return vpploc + vppnl
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 get_pp(cell, kpt=np.zeros(3)): '''Get the periodic pseudotential nuc-el AO matrix ''' from pyscf.pbc import tools coords = cell.get_uniform_grids() aoR = cell.pbc_eval_gto('GTOval', coords, kpt=kpt) nao = cell.nao_nr() SI = cell.get_SI() vlocG = get_vlocG(cell) vpplocG = -np.sum(SI * vlocG, axis=0) vpplocG[0] = np.sum(get_alphas(cell)) # from get_jvloc_G0 function # vpploc evaluated in real-space vpplocR = tools.ifft(vpplocG, cell.mesh).real vpploc = np.dot(aoR.T.conj(), vpplocR.reshape(-1,1)*aoR) # vppnonloc evaluated in reciprocal space aokG = tools.fftk(np.asarray(aoR.T, order='C'), cell.mesh, np.exp(-1j*np.dot(coords, kpt))).T ngrids = len(aokG) fakemol = mole.Mole() fakemol._atm = np.zeros((1,mole.ATM_SLOTS), dtype=np.int32) fakemol._bas = np.zeros((1,mole.BAS_SLOTS), dtype=np.int32) ptr = mole.PTR_ENV_START fakemol._env = np.zeros(ptr+10) fakemol._bas[0,mole.NPRIM_OF ] = 1 fakemol._bas[0,mole.NCTR_OF ] = 1 fakemol._bas[0,mole.PTR_EXP ] = ptr+3 fakemol._bas[0,mole.PTR_COEFF] = ptr+4 Gv = np.asarray(cell.Gv+kpt) G_rad = lib.norm(Gv, axis=1) vppnl = np.zeros((nao,nao), dtype=np.complex128) for ia in range(cell.natm): symb = cell.atom_symbol(ia) if symb not in cell._pseudo: continue pp = cell._pseudo[symb] for l, proj in enumerate(pp[5:]): rl, nl, hl = proj if nl > 0: hl = np.asarray(hl) fakemol._bas[0,mole.ANG_OF] = l fakemol._env[ptr+3] = .5*rl**2 fakemol._env[ptr+4] = rl**(l+1.5)*np.pi**1.25 pYlm_part = fakemol.eval_gto('GTOval', Gv) pYlm = np.empty((nl,l*2+1,ngrids)) for k in range(nl): qkl = _qli(G_rad*rl, l, k) pYlm[k] = pYlm_part.T * qkl # pYlm is real SPG_lmi = np.einsum('g,nmg->nmg', SI[ia].conj(), pYlm) SPG_lm_aoG = np.einsum('nmg,gp->nmp', SPG_lmi, aokG) tmp = np.einsum('ij,jmp->imp', hl, SPG_lm_aoG) vppnl += np.einsum('imp,imq->pq', SPG_lm_aoG.conj(), tmp) vppnl *= (1./ngrids**2) if aoR.dtype == np.double: return vpploc.real + vppnl.real else: return vpploc + vppnl
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) ao1_dtype = np.result_type(*ao1_kpts) ao2_dtype = np.result_type(*ao2_kpts) 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_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 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) 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) 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)
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 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((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) 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('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_for_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_for_G0(cell, kpts, dms, vk_kpts, kpts_band=kpts_band) return _format_jks(vk_kpts, dm_kpts, input_band, kpts)
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) rho = None 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_kpts(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)] 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 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 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 p0, p1 in lib.prange(0, nao, blksize): for i in range(nset): rho1 = numpy.einsum('ig,jg->ijg', ao1T[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): vk_kpts[i,k1] += weight * lib.dot(vR_dm[i], ao1T.T) vk_kpts = mpi.reduce(lib.asarray(vk_kpts)) 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)
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)