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 write_mo(fout, mol, mo_coeff, mo_energy=None, mo_occ=None): nmo = mo_coeff.shape[1] mo_cart = [] centers = [] types = [] exps = [] p0 = 0 for ib in range(mol.nbas): ia = mol.bas_atom(ib) l = mol.bas_angular(ib) es = mol.bas_exp(ib) c = mol.bas_ctr_coeff(ib) np, nc = c.shape nd = nc*(2*l+1) mosub = mo_coeff[p0:p0+nd].reshape(-1,nc,nmo) c2s = gto.cart2sph(l) mosub = numpy.einsum('yki,cy,pk->pci', mosub, c2s, c) mo_cart.append(mosub.reshape(-1,nmo)) for t in TYPE_MAP[l]: types.append([t]*np) ncart = gto.len_cart(l) exps.extend([es]*ncart) centers.extend([ia+1]*(np*ncart)) p0 += nd mo_cart = numpy.vstack(mo_cart) centers = numpy.hstack(centers) types = numpy.hstack(types) exps = numpy.hstack(exps) nprim, nmo = mo_cart.shape fout.write('title line\n') fout.write('GAUSSIAN %3d MOL ORBITALS %3d PRIMITIVES %3d NUCLEI\n' % (mo_cart.shape[1], mo_cart.shape[0], mol.natm)) for ia in range(mol.natm): x, y, z = mol.atom_coord(ia) fout.write('%-4s %-4d (CENTRE%3d) %12.8f %12.8f %12.8f CHARGE = %.1f\n' % (mol.atom_pure_symbol(ia), ia, ia, x, y, z, mol.atom_charge(ia))) for i0, i1 in lib.prange(0, nprim, 20): fout.write('CENTRE ASSIGNMENTS %s\n' % ''.join('%3d'%x for x in centers[i0:i1])) for i0, i1 in lib.prange(0, nprim, 20): fout.write('TYPE ASSIGNMENTS %s\n' % ''.join('%3d'%x for x in types[i0:i1])) for i0, i1 in lib.prange(0, nprim, 5): fout.write('EXPONENTS %s\n' % ' '.join('%14.7E'%x for x in exps[i0:i1])) for k in range(nmo): mo = mo_cart[:,k] if mo_energy is None or mo_occ is None: fout.write('CANMO %d\n' (k, mo_occ[k], mo_energy[k])) else: fout.write('MO %-4d OCC NO = %12.8f ORB. ENERGY = %12.8f\n' % (k, mo_occ[k], mo_energy[k])) for i0, i1 in lib.prange(0, nprim, 5): fout.write('%s\n' % ' '.join('%15.8E'%x for x in mo[i0:i1]))
def pw_loop(self, mol, auxmol, gs, shls_slice=None, max_memory=2000): '''Plane wave part''' if isinstance(gs, int): gs = [gs]*3 naux = auxmol.nao_nr() Gv, Gvbase, kws = non_uniform_kgrids(gs) nxyz = [i*2 for i in gs] gxyz = lib.cartesian_prod([range(i) for i in nxyz]) kk = numpy.einsum('ki,ki->k', Gv, Gv) idx = numpy.argsort(kk)[::-1] # idx = idx[(kk[idx] < 300.) & (kk[idx] > 1e-4)] # ~ Cut high energy plain waves # log.debug('Cut grids %d to %d', Gv.shape[0], len(idx)) kk = kk[idx] Gv = Gv[idx] kws = kws[idx] gxyz = gxyz[idx] coulG = .5/numpy.pi**2 * kws / kk if shls_slice is None: ni = nj = mol.nao_nr() else: ao_loc = mol.ao_loc_nr() ni = ao_loc[shls_slice[1]] - ao_loc[shls_slice[0]] nj = ao_loc[shls_slice[3]] - ao_loc[shls_slice[2]] nij = ni * nj blksize = min(max(16, int(max_memory*1e6*.7/16/nij)), 16384) sublk = max(16,int(blksize//4)) pqkRbuf = numpy.empty(nij*sublk) pqkIbuf = numpy.empty(nij*sublk) LkRbuf = numpy.empty(naux*sublk) LkIbuf = numpy.empty(naux*sublk) for p0, p1 in lib.prange(0, coulG.size, blksize): aoao = ft_ao.ft_aopair(mol, Gv[p0:p1], shls_slice, 's1', Gvbase, gxyz[p0:p1], nxyz) aoaux = ft_ao.ft_ao(auxmol, Gv[p0:p1], None, Gvbase, gxyz[p0:p1], nxyz) for i0, i1 in lib.prange(0, p1-p0, sublk): nG = i1 - i0 pqkR = numpy.ndarray((ni,nj,nG), buffer=pqkRbuf) pqkI = numpy.ndarray((ni,nj,nG), buffer=pqkIbuf) LkR = numpy.ndarray((naux,nG), buffer=LkRbuf) LkI = numpy.ndarray((naux,nG), buffer=LkIbuf) pqkR[:] = aoao[i0:i1].real.transpose(1,2,0) pqkI[:] = aoao[i0:i1].imag.transpose(1,2,0) LkR [:] = aoaux[i0:i1].real.T LkI [:] = aoaux[i0:i1].imag.T yield (pqkR.reshape(-1,nG), LkR, pqkI.reshape(-1,nG), LkI, coulG[p0+i0:p0+i1]) aoao = aoaux = None
def send(sendbuf, dest=0, tag=0): sendbuf = numpy.asarray(sendbuf, order='C') comm.send((sendbuf.shape, sendbuf.dtype), dest=dest, tag=tag) send_seg = numpy.ndarray(sendbuf.size, dtype=sendbuf.dtype, buffer=sendbuf) for p0, p1 in lib.prange(0, sendbuf.size, BLKSIZE): comm.Send(send_seg[p0:p1], dest=dest, tag=tag) return sendbuf
def recv(source=0, tag=0): shape, dtype = comm.recv(source=source, tag=tag) recvbuf = numpy.empty(shape, dtype=dtype) recv_seg = numpy.ndarray(recvbuf.size, dtype=recvbuf.dtype, buffer=recvbuf) for p0, p1 in lib.prange(0, recvbuf.size, BLKSIZE): comm.Recv(recv_seg[p0:p1], source=source, tag=tag) return recvbuf
def get_hcore(self, mol=None): if mol is None: mol = self.mol if getattr(scf_method, 'get_hcore', None): h1e = method_class.get_hcore(self, mol) else: # DO NOT modify post-HF objects to avoid the MM charges applied twice raise RuntimeError('mm_charge function cannot be applied on post-HF methods') if pyscf.DEBUG: v = 0 for i,q in enumerate(charges): mol.set_rinv_origin(coords[i]) v += mol.intor('int1e_rinv') * -q else: if mol.cart: intor = 'int3c2e_cart' else: intor = 'int3c2e_sph' nao = mol.nao max_memory = self.max_memory - lib.current_memory()[0] blksize = int(min(max_memory*1e6/8/nao**2, 200)) v = 0 for i0, i1 in lib.prange(0, charges.size, blksize): fakemol = gto.fakemol_for_charges(coords[i0:i1]) j3c = df.incore.aux_e2(mol, fakemol, intor=intor, aosym='s2ij') v += numpy.einsum('xk,k->x', j3c, -charges[i0:i1]) v = lib.unpack_tril(v) return h1e + v
def init_amps(mycc, eris=None): eris = getattr(mycc, '_eris', None) if eris is None: mycc.ao2mo() eris = mycc._eris time0 = time.clock(), time.time() mo_e = eris.mo_energy nocc = mycc.nocc nvir = mo_e.size - nocc eia = mo_e[:nocc,None] - mo_e[None,nocc:] t1T = eris.fock[nocc:,:nocc] / eia.T loc0, loc1 = _task_location(nvir) t2T = numpy.empty((loc1-loc0,nvir,nocc,nocc)) max_memory = mycc.max_memory - lib.current_memory()[0] blksize = int(min(nvir, max(BLKMIN, max_memory*.3e6/8/(nocc**2*nvir+1)))) emp2 = 0 for p0, p1 in lib.prange(0, loc1-loc0, blksize): eris_ovov = eris.ovov[:,p0:p1] t2T[p0:p1] = (eris_ovov.transpose(1,3,0,2) / lib.direct_sum('ia,jb->abij', eia[:,p0+loc0:p1+loc0], eia)) emp2 += 2 * numpy.einsum('abij,iajb', t2T[p0:p1], eris_ovov) emp2 -= numpy.einsum('abji,iajb', t2T[p0:p1], eris_ovov) mycc.emp2 = comm.allreduce(emp2) logger.info(mycc, 'Init t2, MP2 energy = %.15g', mycc.emp2) logger.timer(mycc, 'init mp2', *time0) return mycc.emp2, t1T.T, t2T.transpose(2,3,0,1)
def get_nuc_less_accurate(mydf, kpts=None): log = logger.Logger(mydf.stdout, mydf.verbose) t1 = t0 = (time.clock(), time.time()) if kpts is None: kpts_lst = numpy.zeros((1, 3)) else: kpts_lst = numpy.reshape(kpts, (-1, 3)) nkpts = len(kpts_lst) if mydf._cderi is None: mydf.build() cell = mydf.cell fused_cell, fuse = fuse_auxcell_(mydf, mydf.auxcell) nao = cell.nao_nr() charge = -cell.atom_charges() j2c = pgto.intor_cross("cint2c2e_sph", fused_cell, _fake_nuc(cell)) jaux = j2c.dot(charge) jaux -= charge.sum() * mydf.auxbar(fused_cell) Gv = cell.get_Gv(mydf.gs) SI = cell.get_SI(Gv) # The normal nuclues have been considered in function get_gth_vlocG_part1 # The result vG is the potential in G-space for erf part of the pp nuclues and # "numpy.dot(charge, SI) * coulG" for normal nuclues. vpplocG = pgto.pseudo.pp_int.get_gth_vlocG_part1(cell, Gv) vG = -1.0 / cell.vol * numpy.einsum("ij,ij->j", SI, vpplocG) kpt_allow = numpy.zeros(3) if is_zero(kpts_lst): vj = numpy.zeros((nkpts, nao ** 2)) else: vj = numpy.zeros((nkpts, nao ** 2), dtype=numpy.complex128) max_memory = max(2000, mydf.max_memory - lib.current_memory()[0]) for k, pqkR, pqkI, p0, p1 in mydf.ft_loop(cell, mydf.gs, kpt_allow, kpts_lst, max_memory=max_memory): if not gamma_point(kpts_lst[k]): vj[k] += numpy.einsum("k,xk->x", vG.real, pqkI) * 1j vj[k] += numpy.einsum("k,xk->x", vG.imag, pqkR) * -1j vj[k] += numpy.einsum("k,xk->x", vG.real, pqkR) vj[k] += numpy.einsum("k,xk->x", vG.imag, pqkI) pqkR = pqkI = None Gv = cell.get_Gv(mydf.gs) aoaux = ft_ao.ft_ao(fused_cell, Gv) jaux -= numpy.einsum("x,xj->j", vG.real, aoaux.real) jaux -= numpy.einsum("x,xj->j", vG.imag, aoaux.imag) jaux = fuse(jaux) vj = vj.reshape(-1, nao, nao) for k, kpt in enumerate(kpts_lst): with mydf.load_Lpq((kpt, kpt)) as Lpq: v = 0 for p0, p1 in lib.prange(0, jaux.size, mydf.blockdim): v += numpy.dot(jaux[p0:p1], numpy.asarray(Lpq[p0:p1])) if gamma_point(kpt): vj[k] += lib.unpack_tril(numpy.asarray(v.real, order="C")) else: vj[k] += lib.unpack_tril(v) if kpts is None or numpy.shape(kpts) == (3,): vj = vj[0] return vj
def write(self, field, fname, comment=None): """ Result: .cube file with the field in the file fname. """ assert(field.ndim == 3) assert(field.shape == (self.nx, self.ny, self.nz)) if comment is None: comment = 'Generic field? Supply the optional argument "comment" to define this line' mol = self.mol coord = mol.atom_coords() with open(fname, 'w') as f: f.write(comment+'\n') f.write('PySCF Version: %s Date: %s\n' % (pyscf.__version__, time.ctime())) f.write('%5d' % mol.natm) f.write('%12.6f%12.6f%12.6f\n' % tuple(self.boxorig.tolist())) f.write('%5d%12.6f%12.6f%12.6f\n' % (self.nx, self.xs[1], 0, 0)) f.write('%5d%12.6f%12.6f%12.6f\n' % (self.ny, 0, self.ys[1], 0)) f.write('%5d%12.6f%12.6f%12.6f\n' % (self.nz, 0, 0, self.zs[1])) for ia in range(mol.natm): chg = mol.atom_charge(ia) f.write('%5d%12.6f'% (chg, chg)) f.write('%12.6f%12.6f%12.6f\n' % tuple(coord[ia])) for ix in range(self.nx): for iy in range(self.ny): for iz0, iz1 in lib.prange(0, self.nz, 6): fmt = '%13.5E' * (iz1-iz0) + '\n' f.write(fmt % tuple(field[ix,iy,iz0:iz1].tolist()))
def extrapolate(self, nd=None): if nd is None: nd = self.get_num_vec() if nd == 0: raise RuntimeError('No vector found in DIIS object.') h = self._H[:nd+1,:nd+1].copy() h[1:,1:] = mpi.comm.allreduce(self._H[1:nd+1,1:nd+1]) g = numpy.zeros(nd+1, h.dtype) g[0] = 1 w, v = scipy.linalg.eigh(h) if numpy.any(abs(w)<1e-14): logger.debug(self, 'Singularity found in DIIS error vector space.') idx = abs(w)>1e-14 c = numpy.dot(v[:,idx]*(1./w[idx]), numpy.dot(v[:,idx].T.conj(), g)) else: try: c = numpy.linalg.solve(h, g) except numpy.linalg.linalg.LinAlgError as e: logger.warn(self, ' diis singular, eigh(h) %s', w) raise e logger.debug1(self, 'diis-c %s', c) xnew = None for i, ci in enumerate(c[1:]): xi = self.get_vec(i) if xnew is None: xnew = numpy.zeros(xi.size, c.dtype) for p0, p1 in lib.prange(0, xi.size, lib.diis.BLOCK_SIZE): xnew[p0:p1] += xi[p0:p1] * ci return xnew
def fuse(Lpq): Lpq, chgLpq = Lpq[:naux], Lpq[naux:] if Lpq.ndim == 1: npq = 1 Lpq_sph = numpy.empty(naux_sph, dtype=Lpq.dtype) else: npq = Lpq.shape[1] Lpq_sph = numpy.empty((naux_sph,npq), dtype=Lpq.dtype) if Lpq.dtype == numpy.complex: npq *= 2 # c2s_fn supports double only, *2 to handle complex for i in range(auxcell.nbas): l = auxcell.bas_angular(i) ia = auxcell.bas_atom(i) p0 = modchg_offset[ia,l] if p0 >= 0: nd = (l+1) * (l+2) // 2 c0, c1 = aux_loc[i], aux_loc[i+1] s0, s1 = aux_loc_sph[i], aux_loc_sph[i+1] for i0, i1 in lib.prange(c0, c1, nd): Lpq[i0:i1] -= chgLpq[p0:p0+nd] if l < 2: Lpq_sph[s0:s1] = Lpq[c0:c1] else: Lpq_cart = numpy.asarray(Lpq[c0:c1], order='C') c2s_fn(Lpq_sph[s0:s1].ctypes.data_as(ctypes.c_void_p), ctypes.c_int(npq * auxcell.bas_nctr(i)), Lpq_cart.ctypes.data_as(ctypes.c_void_p), ctypes.c_int(l)) return Lpq_sph
def __init__(self, mycc): self._cc = mycc self.daemon = None nocc, nvir = mycc.t1.shape nmo = nocc + nvir nvir_seg = (nvir + mpi.pool.size - 1) // mpi.pool.size max_memory = mycc.max_memory - lib.current_memory()[0] max_memory = max(0, max_memory - nocc**3*13*lib.num_threads()*8/1e6) blksize = max(BLKMIN, (max_memory*.5e6/8/(6*nmo*nocc))**.5 - nocc/4) blksize = int(min(comm.allgather(min(nvir/6+2, nvir_seg/2+1, blksize)))) logger.debug1(mycc, 'GlobalDataHandler blksize %s', blksize) self.vranges = [] self.data_partition = [] self.segment_location = {} for task in range(mpi.pool.size): p0 = nvir_seg * task p1 = min(nvir, p0 + nvir_seg) self.vranges.append((p0, p1)) for j0, j1 in lib.prange(p0, p1, blksize): self.data_partition.append((j0, j1)) self.segment_location[j0] = task logger.debug1(mycc, 'data_partition %s', self.data_partition) logger.debug1(mycc, 'segment_location %s', self.segment_location)
def build_Lpq_pbc(mydf, auxcell, kptij_lst): """Fitting coefficients for auxiliary functions""" kpts_ji = kptij_lst[:, 1] - kptij_lst[:, 0] uniq_kpts, uniq_index, uniq_inverse = unique(kpts_ji) max_memory = max(2000, (mydf.max_memory - lib.current_memory()[0])) if mydf.metric.upper() == "S": outcore.aux_e2( mydf.cell, auxcell, mydf._cderi, "cint3c1e_sph", kptij_lst=kptij_lst, dataname="Lpq", max_memory=max_memory ) s_aux = auxcell.pbc_intor("cint1e_ovlp_sph", hermi=1, kpts=uniq_kpts) else: # mydf.metric.upper() == 'T' outcore.aux_e2( mydf.cell, auxcell, mydf._cderi, "cint3c1e_p2_sph", kptij_lst=kptij_lst, dataname="Lpq", max_memory=max_memory, ) s_aux = [x * 2 for x in auxcell.pbc_intor("cint1e_kin_sph", hermi=1, kpts=uniq_kpts)] s_aux = [scipy.linalg.cho_factor(x) for x in s_aux] max_memory = mydf.max_memory - lib.current_memory()[0] naux = auxcell.nao_nr() blksize = max(int(max_memory * 0.5 * 1e6 / 16 / naux / mydf.blockdim), 1) * mydf.blockdim with h5py.File(mydf._cderi) as feri: for k, where in enumerate(uniq_inverse): s_k = s_aux[where] key = "Lpq/%d" % k Lpq = feri[key] nao_pair = Lpq.shape[1] for p0, p1 in lib.prange(0, nao_pair, blksize): Lpq[:, p0:p1] = scipy.linalg.cho_solve(s_k, Lpq[:, p0:p1])
def allgather(sendbuf, split_recvbuf=False): sendbuf = numpy.asarray(sendbuf, order='C') shape = sendbuf.shape attr = comm.allgather((shape, sendbuf.dtype.char)) rshape = [x[0] for x in attr] counts = numpy.array([numpy.prod(x) for x in rshape]) mpi_dtype = numpy.result_type(*[x[1] for x in attr]).char _assert(sendbuf.dtype.char == mpi_dtype or sendbuf.size == 0) displs = numpy.append(0, numpy.cumsum(counts[:-1])) recvbuf = numpy.empty(sum(counts), dtype=mpi_dtype) sendbuf = sendbuf.ravel() for p0, p1 in lib.prange(0, numpy.max(counts), BLKSIZE): counts_seg = _segment_counts(counts, p0, p1) comm.Allgatherv([sendbuf[p0:p1], mpi_dtype], [recvbuf, counts_seg, displs+p0, mpi_dtype]) if split_recvbuf: return [recvbuf[p0:p0+c].reshape(shape) for p0,c,shape in zip(displs,counts,rshape)] else: try: return recvbuf.reshape((-1,) + shape[1:]) except ValueError: return recvbuf
def energy(mycc, t1=None, t2=None, eris=None): '''CCSD correlation energy''' if t1 is None: t1 = mycc.t1 if t2 is None: t2 = mycc.t2 eris = getattr(mycc, '_eris', None) if eris is None: mycc.ao2mo() eris = mycc._eris nocc, nvir = t1.shape t2T = t2.transpose(2,3,0,1) fock = eris.fock loc0, loc1 = _task_location(nvir) e = numpy.einsum('ia,ia', fock[:nocc,nocc:], t1) * 2 max_memory = mycc.max_memory - lib.current_memory()[0] blksize = int(min(nvir, max(BLKMIN, max_memory*.3e6/8/(nocc**2*nvir+1)))) for p0, p1 in lib.prange(0, loc1-loc0, blksize): eris_ovov = eris.ovov[:,p0:p1] tau = t2T[p0:p1] + numpy.einsum('ia,jb->abij', t1[:,p0+loc0:p1+loc0], t1) e += 2 * numpy.einsum('abij,iajb', tau, eris_ovov) e -= numpy.einsum('abji,iajb', tau, eris_ovov) e = comm.allreduce(e) if rank == 0 and abs(e.imag) > 1e-4: logger.warn(mycc, 'Non-zero imaginary part found in CCSD energy %s', e) return e.real
def _make_eris_outcore(mycc, mo_coeff=None): cput0 = (time.clock(), time.time()) log = logger.Logger(mycc.stdout, mycc.verbose) eris = _ChemistsERIs() eris._common_init_(mycc, mo_coeff) mol = mycc.mol mo_coeff = eris.mo_coeff nocc = eris.nocc nao, nmo = mo_coeff.shape nvir = nmo - nocc orbo = mo_coeff[:,:nocc] orbv = mo_coeff[:,nocc:] nvpair = nvir * (nvir+1) // 2 eris.feri1 = lib.H5TmpFile() eris.oooo = eris.feri1.create_dataset('oooo', (nocc,nocc,nocc,nocc), 'f8') eris.ovoo = eris.feri1.create_dataset('ovoo', (nocc,nvir,nocc,nocc), 'f8', chunks=(nocc,1,nocc,nocc)) eris.ovov = eris.feri1.create_dataset('ovov', (nocc,nvir,nocc,nvir), 'f8', chunks=(nocc,1,nocc,nvir)) eris.ovvo = eris.feri1.create_dataset('ovvo', (nocc,nvir,nvir,nocc), 'f8', chunks=(nocc,1,nvir,nocc)) eris.ovvv = eris.feri1.create_dataset('ovvv', (nocc,nvir,nvir,nvir), 'f8') eris.oovv = eris.feri1.create_dataset('oovv', (nocc,nocc,nvir,nvir), 'f8', chunks=(nocc,nocc,1,nvir)) eris.vvvv = eris.feri1.create_dataset('vvvv', (nvir,nvir,nvir,nvir), 'f8') max_memory = max(MEMORYMIN, mycc.max_memory-lib.current_memory()[0]) ftmp = lib.H5TmpFile() ao2mo.full(mol, mo_coeff, ftmp, max_memory=max_memory, verbose=log) eri = ftmp['eri_mo'] nocc_pair = nocc*(nocc+1)//2 tril2sq = lib.square_mat_in_trilu_indices(nmo) oo = eri[:nocc_pair] eris.oooo[:] = ao2mo.restore(1, oo[:,:nocc_pair], nocc) oovv = lib.take_2d(oo, tril2sq[:nocc,:nocc].ravel(), tril2sq[nocc:,nocc:].ravel()) eris.oovv[:] = oovv.reshape(nocc,nocc,nvir,nvir) oo = oovv = None tril2sq = lib.square_mat_in_trilu_indices(nmo) blksize = min(nvir, max(BLKMIN, int(max_memory*1e6/8/nmo**3/2))) for p0, p1 in lib.prange(0, nvir, blksize): q0, q1 = p0+nocc, p1+nocc off0 = q0*(q0+1)//2 off1 = q1*(q1+1)//2 buf = lib.unpack_tril(eri[off0:off1]) tmp = buf[ tril2sq[q0:q1,:nocc] - off0 ] eris.ovoo[:,p0:p1] = tmp[:,:,:nocc,:nocc].transpose(1,0,2,3) eris.ovvo[:,p0:p1] = tmp[:,:,nocc:,:nocc].transpose(1,0,2,3) eris.ovov[:,p0:p1] = tmp[:,:,:nocc,nocc:].transpose(1,0,2,3) eris.ovvv[:,p0:p1] = tmp[:,:,nocc:,nocc:].transpose(1,0,2,3) tmp = buf[ tril2sq[q0:q1,nocc:q1] - off0 ] eris.vvvv[p0:p1,:p1] = tmp[:,:,nocc:,nocc:] if p0 > 0: eris.vvvv[:p0,p0:p1] = tmp[:,:p0,nocc:,nocc:].transpose(1,0,2,3) buf = tmp = None log.timer('CCSD integral transformation', *cput0) return eris
def get_vk_s4(mol, dm): ao_loc = mol.ao_loc_nr() nao = ao_loc[-1] vk = numpy.zeros((nao,nao)) bas_groups = list(lib.prange(0, mol.nbas, 3)) for ip, (ish0, ish1) in enumerate(bas_groups): for jp, (jsh0, jsh1) in enumerate(bas_groups[:ip]): for kp, (ksh0, ksh1) in enumerate(bas_groups): for lp, (lsh0, lsh1) in enumerate(bas_groups[:kp]): shls_slice = (ish0, ish1, jsh0, jsh1, ksh0, ksh1, lsh0, lsh1) i0, i1, j0, j1, k0, k1, l0, l1 = [ao_loc[x] for x in shls_slice] dms = [dm[j0:j1,k0:k1], dm[i0:i1,k0:k1], dm[j0:j1,l0:l1], dm[i0:i1,l0:l1]] scripts = ['ijkl,jk->il', 'ijkl,ik->jl', 'ijkl,jl->ik', 'ijkl,il->jk'] kparts = jk.get_jk(mol, dms, scripts, shls_slice=shls_slice) vk[i0:i1,l0:l1] += kparts[0] vk[j0:j1,l0:l1] += kparts[1] vk[i0:i1,k0:k1] += kparts[2] vk[j0:j1,k0:k1] += kparts[3] lsh0, lsh1 = ksh0, ksh1 shls_slice = (ish0, ish1, jsh0, jsh1, ksh0, ksh1, lsh0, lsh1) i0, i1, j0, j1, k0, k1, l0, l1 = [ao_loc[x] for x in shls_slice] kparts = jk.get_jk(mol, [dm[j0:j1,k0:k1], dm[i0:i1,k0:k1]], scripts=['ijkl,jk->il', 'ijkl,ik->jl'], shls_slice=shls_slice) vk[i0:i1,l0:l1] += kparts[0] vk[j0:j1,l0:l1] += kparts[1] jsh0, jsh1 = ish0, ish1 for kp, (ksh0, ksh1) in enumerate(bas_groups): for lp, (lsh0, lsh1) in enumerate(bas_groups[:kp]): shls_slice = (ish0, ish1, jsh0, jsh1, ksh0, ksh1, lsh0, lsh1) i0, i1, j0, j1, k0, k1, l0, l1 = [ao_loc[x] for x in shls_slice] kparts = jk.get_jk(mol, [dm[j0:j1,k0:k1], dm[j0:j1,l0:l1]], scripts=['ijkl,jk->il', 'ijkl,jl->ik'], shls_slice=shls_slice) vk[i0:i1,l0:l1] += kparts[0] vk[j0:j1,k0:k1] += kparts[1] lsh0, lsh1 = ksh0, ksh1 shls_slice = (ish0, ish1, jsh0, jsh1, ksh0, ksh1, lsh0, lsh1) i0, i1, j0, j1, k0, k1, l0, l1 = [ao_loc[x] for x in shls_slice] kparts = jk.get_jk(mol, [dm[j0:j1,k0:k1]], scripts=['ijkl,jk->il'], shls_slice=shls_slice) vk[i0:i1,l0:l1] += kparts[0] return vk
def gather(sendbuf, root=0, split_recvbuf=False): # if pool.debug: # if rank == 0: # res = [sendbuf] # for k in range(1, pool.size): # dat = comm.recv(source=k) # res.append(dat) # return numpy.vstack([x for x in res if len(x) > 0]) # else: # comm.send(sendbuf, dest=0) # return sendbuf sendbuf = numpy.asarray(sendbuf, order='C') shape = sendbuf.shape size_dtype = comm.allgather((shape, sendbuf.dtype.char)) rshape = [x[0] for x in size_dtype] counts = numpy.array([numpy.prod(x) for x in rshape]) mpi_dtype = numpy.result_type(*[x[1] for x in size_dtype]).char _assert(sendbuf.dtype == mpi_dtype or sendbuf.size == 0) if rank == root: displs = numpy.append(0, numpy.cumsum(counts[:-1])) recvbuf = numpy.empty(sum(counts), dtype=mpi_dtype) sendbuf = sendbuf.ravel() for p0, p1 in lib.prange(0, numpy.max(counts), BLKSIZE): counts_seg = _segment_counts(counts, p0, p1) comm.Gatherv([sendbuf[p0:p1], mpi_dtype], [recvbuf, counts_seg, displs+p0, mpi_dtype], root) if split_recvbuf: return [recvbuf[p0:p0+c].reshape(shape) for p0,c,shape in zip(displs,counts,rshape)] else: try: return recvbuf.reshape((-1,) + shape[1:]) except ValueError: return recvbuf else: send_seg = sendbuf.ravel() for p0, p1 in lib.prange(0, numpy.max(counts), BLKSIZE): comm.Gatherv([send_seg[p0:p1], mpi_dtype], None, root) return sendbuf
def bcast(buf, root=0): buf = numpy.asarray(buf, order='C') shape, dtype = comm.bcast((buf.shape, buf.dtype.char)) if rank != root: buf = numpy.empty(shape, dtype=dtype) buf_seg = numpy.ndarray(buf.size, dtype=buf.dtype, buffer=buf) for p0, p1 in lib.prange(0, buf.size, BLKSIZE): comm.Bcast(buf_seg[p0:p1], root) return buf
def allreduce(sendbuf, op=MPI.SUM): sendbuf = numpy.asarray(sendbuf, order='C') shape, mpi_dtype = comm.bcast((sendbuf.shape, sendbuf.dtype.char)) _assert(sendbuf.shape == shape and sendbuf.dtype.char == mpi_dtype) recvbuf = numpy.zeros_like(sendbuf) send_seg = numpy.ndarray(sendbuf.size, dtype=sendbuf.dtype, buffer=sendbuf) recv_seg = numpy.ndarray(recvbuf.size, dtype=recvbuf.dtype, buffer=recvbuf) for p0, p1 in lib.prange(0, sendbuf.size, BLKSIZE): comm.Allreduce(send_seg[p0:p1], recv_seg[p0:p1], op) return recvbuf
def fuse(Lpq): Lpq, chgLpq = Lpq[:naux], Lpq[naux:] for i in range(auxcell.nbas): l = auxcell.bas_angular(i) ia = auxcell.bas_atom(i) p0 = modchg_offset[ia,l] if p0 >= 0: nd = l * 2 + 1 for i0, i1 in lib.prange(aux_loc[i], aux_loc[i+1], nd): Lpq[i0:i1] -= chgLpq[p0:p0+nd] return Lpq
def loop(self): coulG = tools.get_coulG(self.cell, numpy.zeros(3), gs=mydf.gs) ngs = len(coulG) ao_pairs_G = get_ao_pairs_G(mydf, kptijkl[:2], compact=True) ao_pairs_G *= numpy.sqrt(coulG * (cell.vol / ngs ** 2)).reshape(-1, 1) Lpq = numpy.empty((self.blockdim, ao_pairs_G.shape[1])) for p0, p1 in lib.prange(0, ngs, self.blockdim): Lpq[: p1 - p0] = ao_pairs_G[p0:p1].real yield Lpq[: p1 - p0] Lpq[: p1 - p0] = ao_pairs_G[p0:p1].imag yield Lpq[: p1 - p0]
def prange(self, start, stop, step=None): # affect pw_loop and ft_loop function size = stop - start mpi_size = mpi.pool.size segsize = (size+mpi_size-1) // mpi_size if step is None: step = segsize else: step = min(step, segsize) start = min(size, start + rank * segsize) stop = min(size, start + segsize) return lib.prange(start, stop, step)
def sr_loop(self, kpti_kptj=numpy.zeros((2,3)), max_memory=2000, compact=True, blksize=None): '''Short range part''' kpti, kptj = kpti_kptj unpack = is_zero(kpti-kptj) and not compact is_real = is_zero(kpti_kptj) nao = self.cell.nao_nr() if blksize is None: if is_real: if unpack: blksize = max_memory*1e6/8/(nao*(nao+1)//2+nao**2*2) else: blksize = max_memory*1e6/8/(nao*(nao+1)*2) else: blksize = max_memory*1e6/16/(nao**2*3) blksize = max(16, min(int(blksize), self.blockdim)) logger.debug2(self, 'max_memory %d MB, blksize %d', max_memory, blksize) if unpack: buf = numpy.empty((blksize,nao*(nao+1)//2)) def load(Lpq, b0, b1, bufR, bufI): Lpq = numpy.asarray(Lpq[b0:b1]) if is_real: if unpack: LpqR = lib.unpack_tril(Lpq, out=bufR).reshape(-1,nao**2) else: LpqR = Lpq LpqI = numpy.zeros_like(LpqR) else: shape = Lpq.shape if unpack: tmp = numpy.ndarray(shape, buffer=buf) tmp[:] = Lpq.real LpqR = lib.unpack_tril(tmp, out=bufR).reshape(-1,nao**2) tmp[:] = Lpq.imag LpqI = lib.unpack_tril(tmp, lib.ANTIHERMI, out=bufI).reshape(-1,nao**2) else: LpqR = numpy.ndarray(shape, buffer=bufR) LpqR[:] = Lpq.real LpqI = numpy.ndarray(shape, buffer=bufI) LpqI[:] = Lpq.imag return LpqR, LpqI LpqR = LpqI = j3cR = j3cI = None with self.load_Lpq(kpti_kptj) as Lpq: naux = Lpq.shape[0] with self.load_j3c(kpti_kptj) as j3c: for b0, b1 in lib.prange(0, naux, blksize): LpqR, LpqI = load(Lpq, b0, b1, LpqR, LpqI) j3cR, j3cI = load(j3c, b0, b1, j3cR, j3cI) yield LpqR, LpqI, j3cR, j3cI
def make_phi(pcmobj, dm, r_vdw, ui): mol = pcmobj.mol natm = mol.natm coords_1sph, weights_1sph = make_grids_one_sphere(pcmobj.lebedev_order) ngrid_1sph = coords_1sph.shape[0] if not (isinstance(dm, numpy.ndarray) and dm.ndim == 2): dm = dm[0] + dm[1] tril_dm = lib.pack_tril(dm + dm.T) nao = dm.shape[0] diagidx = numpy.arange(nao) diagidx = diagidx * (diagidx + 1) // 2 + diagidx tril_dm[diagidx] *= .5 atom_coords = mol.atom_coords() atom_charges = mol.atom_charges() extern_point_idx = ui > 0 cav_coords = (atom_coords.reshape(natm, 1, 3) + numpy.einsum('r,gx->rgx', r_vdw, coords_1sph)) v_phi = numpy.empty((natm, ngrid_1sph)) for ia in range(natm): # Note (-) sign is not applied to atom_charges, because (-) is explicitly # included in rhs and L matrix d_rs = atom_coords.reshape(-1, 1, 3) - cav_coords[ia] v_phi[ia] = numpy.einsum('z,zp->p', atom_charges, 1. / lib.norm(d_rs, axis=2)) max_memory = pcmobj.max_memory - lib.current_memory()[0] blksize = int(max(max_memory * 1e6 / 8 / nao**2, 400)) cav_coords = cav_coords[extern_point_idx] v_phi_e = numpy.empty(cav_coords.shape[0]) int3c2e = mol._add_suffix('int3c2e') cintopt = gto.moleintor.make_cintopt(mol._atm, mol._bas, mol._env, int3c2e) for i0, i1 in lib.prange(0, cav_coords.shape[0], blksize): fakemol = gto.fakemol_for_charges(cav_coords[i0:i1]) v_nj = df.incore.aux_e2(mol, fakemol, intor=int3c2e, aosym='s2ij', cintopt=cintopt) v_phi_e[i0:i1] = numpy.einsum('x,xk->k', tril_dm, v_nj) v_phi[extern_point_idx] -= v_phi_e ylm_1sph = numpy.vstack(sph.real_sph_vec(coords_1sph, pcmobj.lmax, True)) phi = -numpy.einsum('n,xn,jn,jn->jx', weights_1sph, ylm_1sph, ui, v_phi) return phi
def pw_contract(istep, sh_range, j3cR, j3cI): bstart, bend, ncol = sh_range if aosym == 's2': shls_slice = (bstart, bend, 0, bend) else: shls_slice = (bstart, bend, 0, cell.nbas) for p0, p1 in lib.prange(0, ngrids, Gblksize): dat = ft_ao._ft_aopair_kpts(cell, Gv[p0:p1], shls_slice, aosym, b, gxyz[p0:p1], Gvbase, kpt, adapted_kptjs, out=buf) nG = p1 - p0 for k, ji in enumerate(adapted_ji_idx): aoao = dat[k].reshape(nG, ncol) pqkR = numpy.ndarray((ncol, nG), buffer=pqkRbuf) pqkI = numpy.ndarray((ncol, nG), buffer=pqkIbuf) pqkR[:] = aoao.real.T pqkI[:] = aoao.imag.T lib.dot(kLR[p0:p1].T, pqkR.T, -1, j3cR[k], 1) lib.dot(kLI[p0:p1].T, pqkI.T, -1, j3cR[k], 1) if not (is_zero(kpt) and gamma_point(adapted_kptjs[k])): lib.dot(kLR[p0:p1].T, pqkI.T, -1, j3cI[k], 1) lib.dot(kLI[p0:p1].T, pqkR.T, 1, j3cI[k], 1) for k, ji in enumerate(adapted_ji_idx): if is_zero(kpt) and gamma_point(adapted_kptjs[k]): v = j3cR[k] else: v = j3cR[k] + j3cI[k] * 1j if j2ctag == 'CD': v = scipy.linalg.solve_triangular(j2c, v, lower=True, overwrite_b=True) feri['j3c/%d/%d' % (ji, istep)] = v else: feri['j3c/%d/%d' % (ji, istep)] = lib.dot(j2c, v) # low-dimension systems if j2c_negative is not None: feri['j3c-/%d/%d' % (ji, istep)] = lib.dot(j2c_negative, v)
def reduce(sendbuf, op=MPI.SUM, root=0): sendbuf = numpy.asarray(sendbuf, order='C') shape, mpi_dtype = comm.bcast((sendbuf.shape, sendbuf.dtype.char)) _assert(sendbuf.shape == shape and sendbuf.dtype.char == mpi_dtype) recvbuf = numpy.zeros_like(sendbuf) send_seg = numpy.ndarray(sendbuf.size, dtype=sendbuf.dtype, buffer=sendbuf) recv_seg = numpy.ndarray(recvbuf.size, dtype=recvbuf.dtype, buffer=recvbuf) for p0, p1 in lib.prange(0, sendbuf.size, BLKSIZE): comm.Reduce(send_seg[p0:p1], recv_seg[p0:p1], op, root) if rank == root: return recvbuf else: return sendbuf
def build_Lpq_1c_approx(mydf, auxcell): get_Lpq = pyscf.df.mdf._make_Lpq_atomic_approx(mydf, mydf.cell, auxcell) max_memory = mydf.max_memory - lib.current_memory()[0] naux = auxcell.nao_nr() nao = mydf.cell.nao_nr() nao_pair = nao * (nao+1) // 2 blksize = max(int(max_memory*.5*1e6/8/naux/mydf.blockdim), 1) * mydf.blockdim with h5py.File(mydf._cderi) as feri: if 'Lpq' in feri: del(feri['Lpq']) chunks = (min(mydf.blockdim,naux), min(mydf.blockdim,nao_pair)) # 512K Lpq = feri.create_dataset('Lpq/0', (naux,nao_pair), 'f8', chunks=chunks) for p0, p1 in lib.prange(0, nao_pair, blksize): v = get_Lpq(None, p0, p1) Lpq[:,p0:p1] = get_Lpq(None, p0, p1)
def pw_loop(self, cell, gs=None, kpti_kptj=None, shls_slice=None, max_memory=2000): '''Plane wave part''' if gs is None: gs = self.gs if kpti_kptj is None: kpti = kptj = numpy.zeros(3) else: kpti, kptj = kpti_kptj nao = cell.nao_nr() gxyz = lib.cartesian_prod((numpy.append(range(gs[0]+1), range(-gs[0],0)), numpy.append(range(gs[1]+1), range(-gs[1],0)), numpy.append(range(gs[2]+1), range(-gs[2],0)))) invh = numpy.linalg.inv(cell._h) Gv = 2*numpy.pi * numpy.dot(gxyz, invh) ngs = gxyz.shape[0] # Theoretically, hermitian symmetry can be also found for kpti == kptj: # f_ji(G) = \int f_ji exp(-iGr) = \int f_ij^* exp(-iGr) = [f_ij(-G)]^* # The hermi operation needs reordering the axis-0. It is inefficient if gamma_point(kpti) and gamma_point(kptj): aosym = 's1hermi' else: aosym = 's1' blksize = min(max(16, int(max_memory*1e6*.75/16/nao**2)), 16384) sublk = max(16, int(blksize//4)) buf = [numpy.zeros(nao*nao*blksize, dtype=numpy.complex128)] pqkRbuf = numpy.empty(nao*nao*sublk) pqkIbuf = numpy.empty(nao*nao*sublk) for p0, p1 in self.prange(0, ngs, blksize): #aoao = ft_ao.ft_aopair(cell, Gv[p0:p1], shls_slice, aosym, invh, # gxyz[p0:p1], gs, (kpti, kptj)) aoao = ft_ao._ft_aopair_kpts(cell, Gv[p0:p1], shls_slice, aosym, invh, gxyz[p0:p1], gs, kptj-kpti, kptj.reshape(1,3), out=buf)[0] for i0, i1 in lib.prange(0, p1-p0, sublk): nG = i1 - i0 pqkR = numpy.ndarray((nao,nao,nG), buffer=pqkRbuf) pqkI = numpy.ndarray((nao,nao,nG), buffer=pqkIbuf) pqkR[:] = aoao[i0:i1].real.transpose(1,2,0) pqkI[:] = aoao[i0:i1].imag.transpose(1,2,0) yield (pqkR.reshape(-1,nG), pqkI.reshape(-1,nG), p0+i0, p0+i1) aoao[:] = 0
def mep(mol, outfile, dm, nx=80, ny=80, nz=80, resolution=RESOLUTION): """Calculates the molecular electrostatic potential (MEP) and write out in cube format. Args: mol : Mole Molecule to calculate the electron density for. outfile : str Name of Cube file to be written. dm : ndarray Density matrix of molecule. Kwargs: nx : int Number of grid point divisions in x direction. Note this is function of the molecule's size; a larger molecule will have a coarser representation than a smaller one for the same value. ny : int Number of grid point divisions in y direction. nz : int Number of grid point divisions in z direction. """ cc = Cube(mol, nx, ny, nz, resolution) coords = cc.get_coords() # Nuclear potential at given points Vnuc = 0 for i in range(mol.natm): r = mol.atom_coord(i) Z = mol.atom_charge(i) rp = r - coords Vnuc += Z / numpy.einsum('xi,xi->x', rp, rp)**.5 # Potential of electron density Vele = numpy.empty_like(Vnuc) for p0, p1 in lib.prange(0, Vele.size, 600): fakemol = gto.fakemol_for_charges(coords[p0:p1]) ints = df.incore.aux_e2(mol, fakemol) Vele[p0:p1] = numpy.einsum('ijp,ij->p', ints, dm) MEP = Vnuc - Vele # MEP at each point MEP = MEP.reshape(nx,ny,nz) # Write the potential cc.write(MEP, outfile, 'Molecular electrostatic potential in real space')
def mep(mol, outfile, dm, nx=80, ny=80, nz=80, resolution=RESOLUTION): """Calculates the molecular electrostatic potential (MEP) and write out in cube format. Args: mol : Mole Molecule to calculate the electron density for. outfile : str Name of Cube file to be written. dm : ndarray Density matrix of molecule. Kwargs: nx : int Number of grid point divisions in x direction. Note this is function of the molecule's size; a larger molecule will have a coarser representation than a smaller one for the same value. ny : int Number of grid point divisions in y direction. nz : int Number of grid point divisions in z direction. """ cc = Cube(mol, nx, ny, nz, resolution) coords = cc.get_coords() # Nuclear potential at given points Vnuc = 0 for i in range(mol.natm): r = mol.atom_coord(i) Z = mol.atom_charge(i) rp = r - coords Vnuc += Z / numpy.einsum('xi,xi->x', rp, rp)**.5 # Potential of electron density Vele = numpy.empty_like(Vnuc) for p0, p1 in lib.prange(0, Vele.size, 600): fakemol = gto.fakemol_for_charges(coords[p0:p1]) ints = df.incore.aux_e2(mol, fakemol) Vele[p0:p1] = numpy.einsum('ijp,ij->p', ints, dm) MEP = Vnuc - Vele # MEP at each point MEP = MEP.reshape(nx, ny, nz) # Write the potential cc.write(MEP, outfile, 'Molecular electrostatic potential in real space')
def _sort_eri(mycc, eris, nocc, nvir, vvop, log): cpu1 = (time.clock(), time.time()) mol = mycc.mol nmo = nocc + nvir if mol.symmetry: orbsym = symm.addons.label_orb_symm(mol, mol.irrep_id, mol.symm_orb, eris.mo_coeff, check=False) orbsym = numpy.asarray(orbsym, dtype=numpy.int32) % 10 else: orbsym = numpy.zeros(nmo, dtype=numpy.int32) o_sorted = _irrep_argsort(orbsym[:nocc]) v_sorted = _irrep_argsort(orbsym[nocc:]) vrank = numpy.argsort(v_sorted) max_memory = max(0, mycc.max_memory - lib.current_memory()[0]) max_memory = min(8000, max_memory * .9) blksize = min(nvir, max(16, int(max_memory * 1e6 / 8 / (nvir * nocc * nmo)))) log.debug1('_sort_eri max_memory %g blksize %d', max_memory, blksize) dtype = vvop.dtype with lib.call_in_background(vvop.__setitem__, sync=not mycc.async_io) as save: bufopv = numpy.empty((nocc, nmo, nvir), dtype=dtype) buf1 = numpy.empty_like(bufopv) buf = numpy.empty((nocc, nvir, nvir), dtype=dtype) for j0, j1 in lib.prange(0, nvir, blksize): ovov = numpy.asarray(eris.ovov[:, j0:j1]) #ovvv = numpy.asarray(eris.ovvv[:,j0:j1]) ovvv = eris.get_ovvv(slice(None), slice(j0, j1)) for j in range(j0, j1): oov = ovov[o_sorted, j - j0] ovv = ovvv[o_sorted, j - j0] #if ovv.ndim == 2: # ovv = lib.unpack_tril(ovv, out=buf) bufopv[:, :nocc, :] = oov[:, o_sorted][:, :, v_sorted].conj() bufopv[:, nocc:, :] = ovv[:, v_sorted][:, :, v_sorted].conj() save(vrank[j], bufopv.transpose(2, 0, 1)) bufopv, buf1 = buf1, bufopv cpu1 = log.timer_debug1('transpose %d:%d' % (j0, j1), *cpu1) return orbsym
def bcast(buf, root=0): if size == 1: return buf is_array = isinstance(buf, np.ndarray) buf = np.asarray(buf, order='C') buf = buf.astype(buf.dtype.char) shape, mpi_dtype = comm.bcast((buf.shape, buf.dtype.char)) if rank != root: buf = np.empty(shape, dtype=mpi_dtype) buf_seg = np.ndarray(buf.size, dtype=buf.dtype, buffer=buf) for p0, p1 in lib.prange(0, buf.size, BLKSIZE): comm.Bcast(buf_seg[p0:p1], root) return buf if is_array else buf.ravel()[0]
def matelem_int3d_coo(self, g, v): """ Compute matrix elements of a potential v given on the 3d grid g using blocks along the grid """ from pyscf import lib bsize = int(min(max(160e6 / (self.norbs * 8.0), 1), g.size)) #print(__name__, bsize, g.size*self.norbs*8) v_matelem = np.zeros((self.norbs, self.norbs)) va = v.reshape(-1) wgts = g.weights if type(g.weights) == np.ndarray else np.repeat( g.weights, g.size) for s, f in lib.prange(0, g.size, bsize): ca2o = self.comp_aos_den( g.coords[s:f]) # compute values of atomic orbitals v_w = (wgts[s:f] * va[s:f]).reshape((f - s, 1)) cb2vo = ca2o * v_w v_matelem += np.dot(ca2o.T, cb2vo) return coo_matrix(v_matelem)
def allreduce(sendbuf, root=0, op=getattr(mpi, 'SUM', None)): if size == 1: return sendbuf is_array = isinstance(sendbuf, np.ndarray) sendbuf = np.asarray(sendbuf, order='C') sendbuf = sendbuf.astype(sendbuf.dtype.char) shape, mpi_dtype = comm.bcast((sendbuf.shape, sendbuf.dtype.char)) assert sendbuf.shape == shape and sendbuf.dtype.char == mpi_dtype recvbuf = np.zeros_like(sendbuf) send_seg = np.ndarray(sendbuf.size, dtype=sendbuf.dtype, buffer=sendbuf) recv_seg = np.ndarray(recvbuf.size, dtype=recvbuf.dtype, buffer=recvbuf) for p0, p1 in lib.prange(0, sendbuf.size, BLKSIZE): comm.Allreduce(send_seg[p0:p1], recv_seg[p0:p1], op) return recvbuf if is_array else recvbuf.ravel()[0]
def orbital(mol, outfile, coeff, nx=80, ny=80, nz=80, resolution=RESOLUTION): """Calculate orbital value on real space grid and write out in cube format. Args: mol : Mole Molecule to calculate the electron density for. outfile : str Name of Cube file to be written. coeff : 1D array coeff coefficient. Kwargs: nx : int Number of grid point divisions in x direction. Note this is function of the molecule's size; a larger molecule will have a coarser representation than a smaller one for the same value. Conflicts to keyword resolution. ny : int Number of grid point divisions in y direction. nz : int Number of grid point divisions in z direction. resolution: float Resolution of the mesh grid in the cube box. If resolution is given in the input, the input nx/ny/nz have no effects. The value of nx/ny/nz will be determined by the resolution and the cube box size. """ cc = Cube(mol, nx, ny, nz, resolution) # Compute density on the .cube grid coords = cc.get_coords() ngrids = cc.get_ngrids() blksize = min(8000, ngrids) orb_on_grid = numpy.empty(ngrids) for ip0, ip1 in lib.prange(0, ngrids, blksize): ao = numint.eval_ao(mol, coords[ip0:ip1]) orb_on_grid[ip0:ip1] = numpy.dot(ao, coeff) orb_on_grid = orb_on_grid.reshape(cc.nx, cc.ny, cc.nz) # Write out orbital to the .cube file cc.write(orb_on_grid, outfile, comment='Orbital value in real space (1/Bohr^3)') return orb_on_grid
def make_phi(pcmobj, dm, r_vdw, ui): mol = pcmobj.mol natm = mol.natm coords_1sph, weights_1sph = make_grids_one_sphere(pcmobj.lebedev_order) ngrid_1sph = coords_1sph.shape[0] if not (isinstance(dm, numpy.ndarray) and dm.ndim == 2): dm = dm[0] + dm[1] tril_dm = lib.pack_tril(dm+dm.T) nao = dm.shape[0] diagidx = numpy.arange(nao) diagidx = diagidx*(diagidx+1)//2 + diagidx tril_dm[diagidx] *= .5 atom_coords = mol.atom_coords() atom_charges = mol.atom_charges() extern_point_idx = ui > 0 cav_coords = (atom_coords.reshape(natm,1,3) + numpy.einsum('r,gx->rgx', r_vdw, coords_1sph)) v_phi = numpy.empty((natm,ngrid_1sph)) for ia in range(natm): # Note (-) sign is not applied to atom_charges, because (-) is explicitly # included in rhs and L matrix d_rs = atom_coords.reshape(-1,1,3) - cav_coords[ia] v_phi[ia] = numpy.einsum('z,zp->p', atom_charges, 1./lib.norm(d_rs,axis=2)) max_memory = pcmobj.max_memory - lib.current_memory()[0] blksize = int(max(max_memory*1e6/8/nao**2, 400)) cav_coords = cav_coords[extern_point_idx] v_phi_e = numpy.empty(cav_coords.shape[0]) int3c2e = mol._add_suffix('int3c2e') for i0, i1 in lib.prange(0, cav_coords.shape[0], blksize): fakemol = gto.fakemol_for_charges(cav_coords[i0:i1]) v_nj = df.incore.aux_e2(mol, fakemol, intor=int3c2e, aosym='s2ij') v_phi_e[i0:i1] = numpy.einsum('x,xk->k', tril_dm, v_nj) v_phi[extern_point_idx] -= v_phi_e ylm_1sph = numpy.vstack(sph.real_sph_vec(coords_1sph, pcmobj.lmax, True)) phi = -numpy.einsum('n,xn,jn,jn->jx', weights_1sph, ylm_1sph, ui, v_phi) return phi
def loop(self, blksize=None): if blksize is None: blksize = self.blockdim kpts0 = numpy.zeros((2, 3)) coulG = tools.get_coulG(self.cell, numpy.zeros(3), mesh=self.mesh, low_dim_ft_type=self.low_dim_ft_type) ngrids = len(coulG) ao_pairs_G = self.get_ao_pairs_G(kpts0, compact=True) ao_pairs_G *= numpy.sqrt(coulG * (self.cell.vol / ngrids**2)).reshape( -1, 1) Lpq = numpy.empty((self.blockdim, ao_pairs_G.shape[1])) for p0, p1 in lib.prange(0, ngrids, blksize): Lpq[:p1 - p0] = ao_pairs_G[p0:p1].real yield Lpq[:p1 - p0] Lpq[:p1 - p0] = ao_pairs_G[p0:p1].imag yield Lpq[:p1 - p0]
def sr_loop(self, kpti_kptj=numpy.zeros((2,3)), max_memory=2000, compact=True): '''Short range part''' kpti, kptj = kpti_kptj unpack = is_zero(kpti-kptj) and not compact nao = self.cell.nao_nr() if is_zero(kpti_kptj) and compact: nao_pair = nao * (nao+1) // 2 else: nao_pair = nao ** 2 if is_zero(kpti_kptj): blksize = max_memory*1e6/8/(nao_pair*5+nao*(nao+1)//2) else: blksize = max_memory*1e6/16/(nao_pair*5+nao*(nao+1)//2) blksize = max(16, min(int(blksize), self.blockdim)) logger.debug2(self, 'max_memory %d MB, blksize %d', max_memory, blksize) if unpack: buf = numpy.empty((blksize,nao*(nao+1)//2)) def load(Lpq, b0, b1, bufR, bufI): Lpq = numpy.asarray(Lpq[b0:b1]) shape = Lpq.shape if unpack: tmp = numpy.ndarray(shape, buffer=buf) tmp[:] = Lpq.real LpqR = lib.unpack_tril(tmp, out=bufR).reshape(-1,nao**2) tmp[:] = Lpq.imag LpqI = lib.unpack_tril(tmp, lib.ANTIHERMI, out=bufI).reshape(-1,nao**2) else: LpqR = numpy.ndarray(shape, buffer=bufR) LpqR[:] = Lpq.real LpqI = numpy.ndarray(shape, buffer=bufI) LpqI[:] = Lpq.imag return LpqR, LpqI LpqR = LpqI = j3cR = j3cI = None with self.load_Lpq(kpti_kptj) as Lpq: naux = Lpq.shape[0] with self.load_j3c(kpti_kptj) as j3c: for b0, b1 in lib.prange(0, naux, blksize): LpqR, LpqI = load(Lpq, b0, b1, LpqR, LpqI) j3cR, j3cI = load(j3c, b0, b1, j3cR, j3cI) yield LpqR, LpqI, j3cR, j3cI
def AO(mol, C, outfile='AO_dir/AOs', nx=80, ny=80, nz=80, resolution=None, save=True): counter = 1 if not os.path.exists(outfile): os.makedirs(outfile) else: exists = True while exists: cache = outfile + '_' + str(counter) exists = os.path.exists(cache) counter += 1 outfile = cache os.makedirs(cache) cc = Cube(mol, nx, ny, nz, resolution) # Compute density on the .cube grid coords = cc.get_coords() ngrids = cc.get_ngrids() #blksize = min(8000, ngrids) blksize = ngrids rho = np.empty(ngrids) for ip0, ip1 in lib.prange(0, ngrids, blksize): ao = numint.eval_ao(mol, coords[ip0:ip1]) x = (coords[:, 0].min(), coords[:, 0].max(), nx) y = (coords[:, 1].min(), coords[:, 1].max(), ny) z = (coords[:, 2].min(), coords[:, 2].max(), nz) if save: np.savetxt(outfile + '/AOs', ao) np.savetxt(outfile + '/C', C) np.savetxt(outfile + '/coords', np.array([y, x, z])) f = open(outfile + '/geom_basis', 'w') f.write(str(mol.atom) + '_') f.write(str(mol.basis)) f.close() #cc.write(np.array([]),outfile + '_geom',comment = 'Geometry') return ao, np.array([y, x, z])
def get_hcore(self, mol=None): if mol is None: mol = self.mol if getattr(method_class, 'get_hcore', None): h1e = method_class.get_hcore(self, mol) else: # DO NOT modify post-HF objects to avoid the MM charges applied twice raise RuntimeError( 'mm_charge function cannot be applied on post-HF methods') coords = self.mm_mol.atom_coords() charges = self.mm_mol.atom_charges() if pyscf.DEBUG: v = 0 for i, q in enumerate(charges): mol.set_rinv_origin(coords[i]) v += mol.intor('int1e_rinv') * -q else: if mol.cart: intor = 'int3c2e_cart' else: intor = 'int3c2e_sph' nao = mol.nao max_memory = self.max_memory - lib.current_memory()[0] blksize = int(min(max_memory * 1e6 / 8 / nao**2, 200)) if max_memory <= 0: blksize = 1 logger.warn( self, 'Memory estimate for reading point charges is negative. ' 'Trying to read point charges one by one.') cintopt = gto.moleintor.make_cintopt(mol._atm, mol._bas, mol._env, intor) v = 0 for i0, i1 in lib.prange(0, charges.size, blksize): fakemol = gto.fakemol_for_charges(coords[i0:i1]) j3c = df.incore.aux_e2(mol, fakemol, intor=intor, aosym='s2ij', cintopt=cintopt) v += numpy.einsum('xk,k->x', j3c, -charges[i0:i1]) v = lib.unpack_tril(v) return h1e + v
def allgather_new(sendbuf, split_recvbuf=False): """ gather that avoids overflow of displs. """ sendbuf = numpy.asarray(sendbuf, order='C') shape = sendbuf.shape attr = comm.allgather((shape, sendbuf.dtype.char)) rshape = [x[0] for x in attr] counts = numpy.array([numpy.prod(x) for x in rshape]) dtype = mpi_dtype = numpy.result_type(*[x[1] for x in attr]).char _assert(sendbuf.dtype.char == mpi_dtype or sendbuf.size == 0) matched = all([x[1:] == rshape[0][1:] for x in rshape[1:]]) if matched: elem_dtype = MPI._typedict.get(mpi_dtype) each_count = numpy.prod(rshape[0][1:]) counts = counts // each_count mpi_dtype = MPI.Datatype(elem_dtype).Create_contiguous( each_count).Commit() else: each_count = 1 displs = numpy.append(0, numpy.cumsum(counts[:-1])) recvbuf = numpy.empty(sum(counts * each_count), dtype=dtype) sendbuf = sendbuf.ravel() for p0, p1 in lib.prange(0, numpy.max(counts), BLKSIZE): counts_seg = _segment_counts(counts, p0, p1) comm.Allgatherv([sendbuf[p0 * each_count:p1 * each_count], mpi_dtype], [recvbuf, counts_seg, displs + p0, mpi_dtype]) if matched: mpi_dtype.Free() if split_recvbuf: return [ recvbuf[p0 * each_count:(p0 + c) * each_count].reshape(shape) for p0, c, shape in zip(displs, counts, rshape) ] else: try: return recvbuf.reshape((-1, ) + shape[1:]) except ValueError: return recvbuf
def loop(self, blksize=None): if self.cell.dimension < 3: raise RuntimeError('ERIs of 1D and 2D systems are not positive ' 'definite. Current API only supports postive ' 'definite ERIs.') if blksize is None: blksize = self.blockdim kpts0 = numpy.zeros((2,3)) coulG = tools.get_coulG(self.cell, numpy.zeros(3), mesh=self.mesh) ngrids = len(coulG) ao_pairs_G = self.get_ao_pairs_G(kpts0, compact=True) ao_pairs_G *= numpy.sqrt(coulG*(self.cell.vol/ngrids**2)).reshape(-1,1) Lpq = numpy.empty((blksize, ao_pairs_G.shape[1])) for p0, p1 in lib.prange(0, ngrids, blksize): Lpq[:p1-p0] = ao_pairs_G[p0:p1].real yield Lpq[:p1-p0] Lpq[:p1-p0] = ao_pairs_G[p0:p1].imag yield Lpq[:p1-p0]
def build_Lpq_pbc(mydf, auxcell, kptij_lst): '''Fitting coefficients for auxiliary functions''' kpts_ji = kptij_lst[:, 1] - kptij_lst[:, 0] uniq_kpts, uniq_index, uniq_inverse = unique(kpts_ji) max_memory = max(2000, (mydf.max_memory - lib.current_memory()[0])) if mydf.metric.upper() == 'S': outcore.aux_e2(mydf.cell, auxcell, mydf._cderi, 'cint3c1e_sph', kptij_lst=kptij_lst, dataname='Lpq', max_memory=max_memory) s_aux = auxcell.pbc_intor('cint1e_ovlp_sph', hermi=1, kpts=uniq_kpts) else: # mydf.metric.upper() == 'T' outcore.aux_e2(mydf.cell, auxcell, mydf._cderi, 'cint3c1e_p2_sph', kptij_lst=kptij_lst, dataname='Lpq', max_memory=max_memory) s_aux = [ x * 2 for x in auxcell.pbc_intor( 'cint1e_kin_sph', hermi=1, kpts=uniq_kpts) ] s_aux = [scipy.linalg.cho_factor(x) for x in s_aux] max_memory = mydf.max_memory - lib.current_memory()[0] naux = auxcell.nao_nr() blksize = max(int(max_memory * .5 * 1e6 / 16 / naux / mydf.blockdim), 1) * mydf.blockdim with h5py.File(mydf._cderi) as feri: for k, where in enumerate(uniq_inverse): s_k = s_aux[where] key = 'Lpq/%d' % k Lpq = feri[key] nao_pair = Lpq.shape[1] for p0, p1 in lib.prange(0, nao_pair, blksize): Lpq[:, p0:p1] = scipy.linalg.cho_solve(s_k, Lpq[:, p0:p1])
def sr_loop(self, kpti_kptj=numpy.zeros((2,3)), max_memory=2000, compact=True, transpose102=False): '''Short range part''' kpti, kptj = kpti_kptj unpack = is_zero(kpti-kptj) and not compact nao = self.cell.nao_nr() max_memory = max(2000, self.max_memory-lib.current_memory()[0]) if is_zero(kpti_kptj): nao_pair = nao * (nao+1) // 2 blksize = max(16, min(int(max_memory*1e6/8/nao_pair/2), self.blockdim)) else: blksize = max(16, min(int(max_memory*1e6/16/nao**2/2), self.blockdim)) if unpack: buf = numpy.empty((blksize,nao*(nao+1)//2)) def load(Lpq, bufR, bufI): shape = Lpq.shape if unpack: tmp = numpy.ndarray(shape, buffer=buf) tmp[:] = Lpq.real LpqR = lib.unpack_tril(tmp, out=bufR).reshape(-1,nao**2) tmp[:] = Lpq.imag LpqI = lib.unpack_tril(tmp, lib.ANTIHERMI, out=bufI).reshape(-1,nao**2) else: LpqR = numpy.ndarray(shape, buffer=bufR) LpqR[:] = Lpq.real LpqI = numpy.ndarray(shape, buffer=bufI) LpqI[:] = Lpq.imag if transpose102: LpqR = numpy.asarray(LpqR.reshape(-1,nao,nao).transpose(1,0,2), order='C') LpqI = numpy.asarray(LpqI.reshape(-1,nao,nao).transpose(1,0,2), order='C') return LpqR, LpqI LpqR = LpqI = j3cR = j3cI = None with self.load_j3c(kpti_kptj) as j3c: with self.load_Lpq(kpti_kptj) as Lpq: naoaux = j3c.shape[0] for b0, b1 in lib.prange(0, naoaux, blksize): LpqR, LpqI = load(numpy.asarray(Lpq[b0:b1]), LpqR, LpqI) j3cR, j3cI = load(numpy.asarray(j3c[b0:b1]), j3cR, j3cI) yield LpqR, LpqI, j3cR, j3cI
def _sort_eri(mycc, eris, nocc, nvir, vvop, log): cpu1 = (time.clock(), time.time()) mol = mycc.mol nmo = nocc + nvir if mol.symmetry: orbsym = symm.addons.label_orb_symm(mol, mol.irrep_id, mol.symm_orb, eris.mo_coeff, check=False) orbsym = numpy.asarray(orbsym, dtype=numpy.int32) % 10 else: orbsym = numpy.zeros(nmo, dtype=numpy.int32) o_sorted = _irrep_argsort(orbsym[:nocc]) v_sorted = _irrep_argsort(orbsym[nocc:]) vrank = numpy.argsort(v_sorted) max_memory = max(0, mycc.max_memory - lib.current_memory()[0]) max_memory = min(8000, max_memory*.9) blksize = min(nvir, max(16, int(max_memory*1e6/8/(nvir*nocc*nmo)))) log.debug1('_sort_eri max_memory %g blksize %d', max_memory, blksize) dtype = vvop.dtype with lib.call_in_background(vvop.__setitem__, sync=not mycc.async_io) as save: bufopv = numpy.empty((nocc,nmo,nvir), dtype=dtype) buf1 = numpy.empty_like(bufopv) buf = numpy.empty((nocc,nvir,nvir), dtype=dtype) for j0, j1 in lib.prange(0, nvir, blksize): ovov = numpy.asarray(eris.ovov[:,j0:j1]) #ovvv = numpy.asarray(eris.ovvv[:,j0:j1]) ovvv = eris.get_ovvv(slice(None), slice(j0,j1)) for j in range(j0,j1): oov = ovov[o_sorted,j-j0] ovv = ovvv[o_sorted,j-j0] #if ovv.ndim == 2: # ovv = lib.unpack_tril(ovv, out=buf) bufopv[:,:nocc,:] = oov[:,o_sorted][:,:,v_sorted].conj() bufopv[:,nocc:,:] = ovv[:,v_sorted][:,:,v_sorted].conj() save(vrank[j], bufopv.transpose(2,0,1)) bufopv, buf1 = buf1, bufopv cpu1 = log.timer_debug1('transpose %d:%d'%(j0,j1), *cpu1) return orbsym
def energy(mycc, t1=None, t2=None, eris=None): '''CCSD correlation energy''' if t1 is None: t1 = mycc.t1 if t2 is None: t2 = mycc.t2 if eris is None: eris = mycc.ao2mo() nocc, nvir = t1.shape fock = eris.fock e = numpy.einsum('ia,ia', fock[:nocc, nocc:], t1) * 2 max_memory = mycc.max_memory - lib.current_memory()[0] blksize = int( min(nvir, max(BLKMIN, max_memory * .3e6 / 8 / (nocc**2 * nvir + 1)))) for p0, p1 in lib.prange(0, nvir, blksize): eris_ovvo = eris.ovvo[:, p0:p1] tau = t2[:, :, p0:p1] e += 2 * numpy.einsum('ijab,iabj', tau, eris_ovvo) e -= numpy.einsum('jiab,iabj', tau, eris_ovvo) if abs(e.imag) > 1e-4: logger.warn(mycc, 'Non-zero imaginary part found in QCISD energy %s', e) return e.real
def write_coeff(fout, mol, mo_coeff): #fout.write('ALDET\n') nmo = mo_coeff.shape[1] mo_cart = [] centers = [] types = [] exps = [] p0 = 0 for ib in range(mol.nbas): ia = mol.bas_atom(ib) l = mol.bas_angular(ib) es = mol.bas_exp(ib) c = mol._libcint_ctr_coeff(ib) #c = mol.bas_ctr_coeff(ib) np, nc = c.shape nd = nc*(2*l+1) mosub = mo_coeff[p0:p0+nd].reshape(-1,nc,nmo) c2s = gto.cart2sph(l) mosub = numpy.einsum('yki,cy,pk->pci', mosub, c2s, c) mo_cart.append(mosub.transpose(1,0,2).reshape(-1,nmo)) for t in TYPE_MAP[l]: types.append([t]*np) ncart = gto.len_cart(l) exps.extend([es]*ncart) centers.extend([ia+1]*(np*ncart)) p0 += nd mo_cart = numpy.vstack(mo_cart) centers = numpy.hstack(centers) types = numpy.hstack(types) exps = numpy.hstack(exps) nprim, nmo = mo_cart.shape fout.write('From PySCF\n') for k in range(nmo): mo = mo_cart[:,k] fout.write('CANMO %d\n' % (k+1)) for i0, i1 in lib.prange(0, nprim, 5): fout.write(' %s\n' % ' '.join('%15.8E'%x for x in mo[i0:i1])) fout.write('END DATA\n')
def _mm_pot(mol: gto.Mole, mm_mol: gto.Mole) -> np.ndarray: """ this function returns the full mm potential (adapted from: qmmm/itrf.py:get_hcore() in PySCF) """ # settings coords = mm_mol.atom_coords() charges = mm_mol.atom_charges() blksize = BLKSIZE # integrals intor = 'int3c2e_cart' if mol.cart else 'int3c2e_sph' cintopt = gto.moleintor.make_cintopt(mol._atm, mol._bas, mol._env, intor) # compute interaction potential mm_pot = 0 for i0, i1 in lib.prange(0, charges.size, blksize): fakemol = gto.fakemol_for_charges(coords[i0:i1]) j3c = df.incore.aux_e2(mol, fakemol, intor=intor, aosym='s2ij', cintopt=cintopt) mm_pot += np.einsum('xk,k->x', j3c, -charges[i0:i1]) mm_pot = lib.unpack_tril(mm_pot) return mm_pot
def write(self, fname, **kw): import time import pyscf from pyscf import lib """ Result: .cube file with the field in the file fname. """ if 'mol' in kw: mol = kw['mol'] # Obligatory argument coord = mol.atom_coords() zz = [mol.atom_charge(ia) for ia in range(mol.natm)] natm = mol.natm else: zz = kw['a2z'] natm = len(zz) coord = kw['a2xyz'] field = kw['field'] # Obligatory argument? comment = kw['comment'] if 'comment' in kw else 'none' with open(fname, 'w') as f: f.write(comment + '\n') f.write('PySCF Version: {:s} Date: {:s}\n'.format( pyscf.__version__, time.ctime())) bo = self.origin - self.get_mesh_center() f.write(('{:5d}' + '{:12.6f}' * 3 + '\n').format(natm, *bo)) for s, rr in zip(self.shape, self.rr): f.write(('{:5d}' + '{:12.6f}' * 3 + '\n').format(s, *rr[1])) for chg, xyz in zip(zz, coord): f.write( ('{:5d}' + '{:12.6f}' * 4 + '\n').format(chg, chg, *xyz)) for ix in range(self.shape[0]): for iy in range(self.shape[1]): for iz0, iz1 in lib.prange(0, self.shape[2], 6): f.write(('{:13.5e}' * (iz1 - iz0) + '\n').format(*field[ix, iy, iz0:iz1]))
def get_gridss(mol, level=1, gthrd=1e-10): Ktime = (time.clock(), time.time()) grids = dft.gen_grid.Grids(mol) grids.level = level grids.build() ngrids = grids.weights.size mask = [] for p0, p1 in lib.prange(0, ngrids, 10000): ao_v = mol.eval_gto('GTOval', grids.coords[p0:p1]) ao_v *= grids.weights[p0:p1,None] wao_v0 = ao_v mask.append(numpy.any(wao_v0>gthrd, axis=1) | numpy.any(wao_v0<-gthrd, axis=1)) mask = numpy.hstack(mask) grids.coords = grids.coords[mask] grids.weights = grids.weights[mask] logger.debug(mol, 'threshold for grids screening %g', gthrd) logger.debug(mol, 'number of grids %d', grids.weights.size) logger.timer_debug1(mol, "Xg screening", *Ktime) return grids
def write(self, field, fname, comment=None): """ Result: .cube file with the field in the file fname. """ assert (field.ndim == 3) assert (field.shape == (self.nx, self.ny, self.nz)) if comment is None: comment = 'Generic field? Supply the optional argument "comment" to define this line' mol = self.mol coord = mol.atom_coords() with open(fname, 'w') as f: f.write(comment + '\n') f.write( f'PySCF Version: {pyscf.__version__} Date: {time.ctime()}\n') f.write(f'{mol.natm:5d}') f.write('%12.6f%12.6f%12.6f\n' % tuple(self.boxorig.tolist())) dx = self.xs[-1] if len(self.xs) == 1 else self.xs[1] dy = self.ys[-1] if len(self.ys) == 1 else self.ys[1] dz = self.zs[-1] if len(self.zs) == 1 else self.zs[1] delta = (self.box.T * [dx, dy, dz]).T f.write( f'{self.nx:5d}{delta[0,0]:12.6f}{delta[0,1]:12.6f}{delta[0,2]:12.6f}\n' ) f.write( f'{self.ny:5d}{delta[1,0]:12.6f}{delta[1,1]:12.6f}{delta[1,2]:12.6f}\n' ) f.write( f'{self.nz:5d}{delta[2,0]:12.6f}{delta[2,1]:12.6f}{delta[2,2]:12.6f}\n' ) for ia in range(mol.natm): atmsymb = mol.atom_symbol(ia) f.write('%5d%12.6f' % (gto.charge(atmsymb), 0.)) f.write('%12.6f%12.6f%12.6f\n' % tuple(coord[ia])) for ix in range(self.nx): for iy in range(self.ny): for iz0, iz1 in lib.prange(0, self.nz, 6): fmt = '%13.5E' * (iz1 - iz0) + '\n' f.write(fmt % tuple(field[ix, iy, iz0:iz1].tolist()))
def orbital(mol, outfile, coeff, nx=80, ny=80, nz=80, resolution=RESOLUTION): """Calculate orbital value on real space grid and write out in cube format. Args: mol : Mole Molecule to calculate the electron density for. outfile : str Name of Cube file to be written. coeff : 1D array coeff coefficient. Kwargs: nx : int Number of grid point divisions in x direction. Note this is function of the molecule's size; a larger molecule will have a coarser representation than a smaller one for the same value. ny : int Number of grid point divisions in y direction. nz : int Number of grid point divisions in z direction. """ cc = Cube(mol, nx, ny, nz, resolution) # Compute density on the .cube grid coords = cc.get_coords() ngrids = cc.get_ngrids() blksize = min(8000, ngrids) orb_on_grid = numpy.zeros(ngrids, dtype=numpy.complex128) for ip0, ip1 in lib.prange(0, ngrids, blksize): aoa, aob = eval_ao(mol, coords[ip0:ip1], deriv=0) orb_on_grid[ip0:ip1] = numpy.dot(aoa, coeff) orb_on_grid[ip0:ip1] += numpy.dot(aob, coeff) orb_on_grid = orb_on_grid.reshape(cc.nx,cc.ny,cc.nz) # Write out orbital to the .cube file cc.write(orb_on_grid, outfile, comment='Orbital value in real space (1/Bohr^3)')
def density(mol, outfile, dm, nx=80, ny=80, nz=80, resolution=RESOLUTION): """Calculates electron density and write out in cube format. Args: mol : Mole Molecule to calculate the electron density for. outfile : str Name of Cube file to be written. dm : ndarray Density matrix of molecule. Kwargs: nx : int Number of grid point divisions in x direction. Note this is function of the molecule's size; a larger molecule will have a coarser representation than a smaller one for the same value. ny : int Number of grid point divisions in y direction. nz : int Number of grid point divisions in z direction. """ cc = Cube(mol, nx, ny, nz, resolution) # Compute density on the .cube grid coords = cc.get_coords() ngrids = cc.get_ngrids() blksize = min(8000, ngrids) rho = numpy.empty(ngrids) for ip0, ip1 in lib.prange(0, ngrids, blksize): ao = numint.eval_ao(mol, coords[ip0:ip1]) rho[ip0:ip1] = numint.eval_rho(mol, ao, dm) rho = rho.reshape(cc.nx, cc.ny, cc.nz) # Write out density to the .cube file cc.write(rho, outfile, comment='Electron density in real space (e/Bohr^3)')
def init_amps(mycc, eris=None): eris = getattr(mycc, '_eris', None) if eris is None: mycc.ao2mo() eris = mycc._eris time0 = logger.process_clock(), logger.perf_counter() mo_e = eris.mo_energy nocc = mycc.nocc nvir = mo_e.size - nocc eia = mo_e[:nocc, None] - mo_e[None, nocc:] #t1T = eris.fock[nocc:, :nocc] / eia.T t1T = np.zeros_like(eris.fock[nocc:, :nocc]) loc0, loc1 = _task_location(nvir) t2T = np.empty((loc1 - loc0, nvir, nocc, nocc)) max_memory = mycc.max_memory - lib.current_memory()[0] blksize = int( min(nvir, max(BLKMIN, max_memory * .3e6 / 8 / (nocc**2 * nvir + 1)))) emp2 = 0 for p0, p1 in lib.prange(0, loc1 - loc0, blksize): eris_vvoo = eris.xvoo[p0:p1] t2T[p0:p1] = ( eris_vvoo / lib.direct_sum('ia, jb -> abij', eia[:, loc0 + p0:loc0 + p1], eia)) emp2 += np.einsum('abij, abij', t2T[p0:p1], eris_vvoo.conj(), optimize=True).real eris_vvoo = None mycc.emp2 = comm.allreduce(emp2) * 0.25 logger.info(mycc, 'Init t2, MP2 energy = %.15g', mycc.emp2) logger.timer(mycc, 'init mp2', *time0) mycc.t1 = t1T.T mycc.t2 = t2T.transpose(2, 3, 0, 1) return mycc.emp2, mycc.t1, mycc.t2
def dump_normal_mode(mol, results): dump = mol.stdout.write freq_wn = results['freq_wavenumber'] idx = freq_wn.real > 0 freq_wn = freq_wn.real[idx] nfreq = freq_wn.size r_mass = results['reduced_mass'].real[idx] force = results['force_const_dyne'].real[idx] vib_t = results['vib_temperature'].real[idx] mode = results['norm_mode'].real[idx] symbols = [mol.atom_symbol(i) for i in range(mol.natm)] def inline(q, col0, col1): return ''.join('%20.4f' % q[i] for i in range(col0, col1)) def mode_inline(row, col0, col1): return ' '.join('%6.2f%6.2f%6.2f' % (mode[i, row, 0], mode[i, row, 1], mode[i, row, 2]) for i in range(col0, col1)) for col0, col1 in lib.prange(0, nfreq, 3): dump('Mode %s\n' % ''.join('%20d' % i for i in range(col0, col1))) dump('Irrep\n') dump('Freq [cm^-1] %s\n' % inline(freq_wn, col0, col1)) dump('Reduced mass [au] %s\n' % inline(r_mass, col0, col1)) dump('Force const [Dyne/A] %s\n' % inline(force, col0, col1)) dump('Char temp [K] %s\n' % inline(vib_t, col0, col1)) #dump('IR\n') #dump('Raman\n') dump('Normal mode %s\n' % (' x y z' * (col1 - col0))) for j, at in enumerate(symbols): dump(' %4d%4s %s\n' % (j, at, mode_inline(j, col0, col1)))
def get_hcore(self, mol=None): ''' (QM 1e grad) + <-d/dX i|q_mm/r_mm|j>''' if mol is None: mol = self.mol g_qm = scf_grad.get_hcore(mol) nao = g_qm.shape[1] if pyscf.DEBUG: v = 0 for i,q in enumerate(charges): mol.set_rinv_origin(coords[i]) v += mol.intor('int1e_iprinv', comp=3) * q else: if mol.cart: intor = 'int3c2e_ip1_cart' else: intor = 'int3c2e_ip1_sph' nao = mol.nao max_memory = self.max_memory - lib.current_memory()[0] blksize = int(max(max_memory*1e6/8/nao**2, 400)) v = 0 for i0, i1 in lib.prange(0, charges.size, blksize): fakemol = gto.fakemol_for_charges(coords[i0:i1]) j3c = df.incore.aux_e2(mol, fakemol, intor, aosym='s1', comp=3) v += numpy.einsum('ipqk,k->ipq', j3c, charges[i0:i1]) return g_qm + v
def _assemble(mydf, kptij_lst, j3c_jobs, gen_int3c, ft_fuse, cderi_file, fswap, log): t1 = (time.clock(), time.time()) cell = mydf.cell ao_loc = cell.ao_loc_nr() nao = ao_loc[-1] kptis = kptij_lst[:, 0] kptjs = kptij_lst[:, 1] kpt_ji = kptjs - kptis uniq_kpts, uniq_index, uniq_inverse = unique(kpt_ji) aosym_s2 = numpy.einsum('ix->i', abs(kptis - kptjs)) < 1e-9 t2 = t1 j3c_workers = numpy.zeros(len(j3c_jobs), dtype=int) #for job_id, ish0, ish1 in mpi.work_share_partition(j3c_jobs): for job_id, ish0, ish1 in mpi.work_stealing_partition(j3c_jobs): gen_int3c(job_id, ish0, ish1) t2 = log.alltimer_debug2('int j3c %d' % job_id, *t2) for k, kpt in enumerate(uniq_kpts): ft_fuse(job_id, k, ish0, ish1) t2 = log.alltimer_debug2('ft-fuse %d k %d' % (job_id, k), *t2) j3c_workers[job_id] = rank j3c_workers = mpi.allreduce(j3c_workers) log.debug2('j3c_workers %s', j3c_workers) t1 = log.timer_debug1('int3c and fuse', *t1) # Pass 2 # Transpose 3-index tensor and save data in cderi_file feri = h5py.File(cderi_file, 'w') nauxs = [fswap['j2c/%d' % k].shape[0] for k, kpt in enumerate(uniq_kpts)] segsize = (max(nauxs) + mpi.pool.size - 1) // mpi.pool.size naux0 = rank * segsize for k, kptij in enumerate(kptij_lst): naux1 = min(nauxs[uniq_inverse[k]], naux0 + segsize) nrow = max(0, naux1 - naux0) if gamma_point(kptij): dtype = 'f8' else: dtype = 'c16' if aosym_s2[k]: nao_pair = nao * (nao + 1) // 2 else: nao_pair = nao * nao feri.create_dataset('j3c/%d' % k, (nrow, nao_pair), dtype, maxshape=(None, nao_pair)) def get_segs_loc(aosym): off0 = numpy.asarray([ao_loc[i0] for x, i0, i1 in j3c_jobs]) off1 = numpy.asarray([ao_loc[i1] for x, i0, i1 in j3c_jobs]) if aosym: # s2 dims = off1 * (off1 + 1) // 2 - off0 * (off0 + 1) // 2 else: dims = (off1 - off0) * nao #dims = numpy.asarray([ao_loc[i1]-ao_loc[i0] for x,i0,i1 in j3c_jobs]) dims = numpy.hstack( [dims[j3c_workers == w] for w in range(mpi.pool.size)]) job_idx = numpy.hstack( [numpy.where(j3c_workers == w)[0] for w in range(mpi.pool.size)]) segs_loc = numpy.append(0, numpy.cumsum(dims)) segs_loc = [(segs_loc[j], segs_loc[j + 1]) for j in numpy.argsort(job_idx)] return segs_loc segs_loc_s1 = get_segs_loc(False) segs_loc_s2 = get_segs_loc(True) job_ids = numpy.where(rank == j3c_workers)[0] def load(k, p0, p1): naux1 = nauxs[uniq_inverse[k]] slices = [(min(i * segsize + p0, naux1), min(i * segsize + p1, naux1)) for i in range(mpi.pool.size)] segs = [] for p0, p1 in slices: val = [ fswap['j3c-chunks/%d/%d' % (job, k)][p0:p1].ravel() for job in job_ids ] if val: segs.append(numpy.hstack(val)) else: segs.append(numpy.zeros(0)) return segs def save(k, p0, p1, segs): segs = mpi.alltoall(segs) naux1 = nauxs[uniq_inverse[k]] loc0, loc1 = min(p0, naux1 - naux0), min(p1, naux1 - naux0) nL = loc1 - loc0 if nL > 0: if aosym_s2[k]: segs = numpy.hstack([ segs[i0 * nL:i1 * nL].reshape(nL, -1) for i0, i1 in segs_loc_s2 ]) else: segs = numpy.hstack([ segs[i0 * nL:i1 * nL].reshape(nL, -1) for i0, i1 in segs_loc_s1 ]) feri['j3c/%d' % k][loc0:loc1] = segs mem_now = max(comm.allgather(lib.current_memory()[0])) max_memory = max(2000, min(8000, mydf.max_memory - mem_now)) if numpy.all(aosym_s2): if gamma_point(kptij_lst): blksize = max(16, int(max_memory * .5e6 / 8 / nao**2)) else: blksize = max(16, int(max_memory * .5e6 / 16 / nao**2)) else: blksize = max(16, int(max_memory * .5e6 / 16 / nao**2 / 2)) log.debug1('max_momory %d MB (%d in use), blksize %d', max_memory, mem_now, blksize) t2 = t1 with lib.call_in_background(save) as async_write: for k, kptji in enumerate(kptij_lst): for p0, p1 in lib.prange(0, segsize, blksize): segs = load(k, p0, p1) async_write(k, p0, p1, segs) t2 = log.timer_debug1( 'assemble k=%d %d:%d (in %d)' % (k, p0, p1, segsize), *t2) if 'j2c-' in fswap: j2c_kpts_lists = [] for k, kpt in enumerate(uniq_kpts): if ('j2c-/%d' % k) in fswap: adapted_ji_idx = numpy.where(uniq_inverse == k)[0] j2c_kpts_lists.append(adapted_ji_idx) for k in numpy.hstack(j2c_kpts_lists): val = [ numpy.asarray(fswap['j3c-/%d/%d' % (job, k)]).ravel() for job in job_ids ] val = mpi.gather(numpy.hstack(val)) if rank == 0: naux1 = fswap['j3c-/0/%d' % k].shape[0] if aosym_s2[k]: v = [ val[i0 * naux1:i1 * naux1].reshape(naux1, -1) for i0, i1 in segs_loc_s2 ] else: v = [ val[i0 * naux1:i1 * naux1].reshape(naux1, -1) for i0, i1 in segs_loc_s1 ] feri['j3c-/%d' % k] = numpy.hstack(v) if 'j3c-kptij' in feri: del (feri['j3c-kptij']) feri['j3c-kptij'] = kptij_lst t1 = log.alltimer_debug1('assembling j3c', *t1) feri.close()
def ft_fuse(job_id, uniq_kptji_id, sh0, sh1): kpt = uniq_kpts[uniq_kptji_id] # kpt = kptj - kpti adapted_ji_idx = numpy.where(uniq_inverse == uniq_kptji_id)[0] adapted_kptjs = kptjs[adapted_ji_idx] nkptj = len(adapted_kptjs) j2c = numpy.asarray(fswap['j2c/%d' % uniq_kptji_id]) j2ctag = j2ctags[uniq_kptji_id] naux0 = j2c.shape[0] if ('j2c-/%d' % uniq_kptji_id) in fswap: j2c_negative = numpy.asarray(fswap['j2c-/%d' % uniq_kptji_id]) else: j2c_negative = None if is_zero(kpt): aosym = 's2' else: aosym = 's1' if aosym == 's2' and cell.dimension == 3: vbar = fuse(mydf.auxbar(fused_cell)) ovlp = cell.pbc_intor('int1e_ovlp', hermi=1, kpts=adapted_kptjs) ovlp = [lib.pack_tril(s) for s in ovlp] j3cR = [None] * nkptj j3cI = [None] * nkptj i0 = ao_loc[sh0] i1 = ao_loc[sh1] for k, idx in enumerate(adapted_ji_idx): key = 'j3c-chunks/%d/%d' % (job_id, idx) v = numpy.asarray(fswap[key]) if aosym == 's2' and cell.dimension == 3: for i in numpy.where(vbar != 0)[0]: v[i] -= vbar[i] * ovlp[k][i0 * (i0 + 1) // 2:i1 * (i1 + 1) // 2].ravel() j3cR[k] = numpy.asarray(v.real, order='C') if v.dtype == numpy.complex128: j3cI[k] = numpy.asarray(v.imag, order='C') v = None ncol = j3cR[0].shape[1] Gblksize = max(16, int(max_memory * 1e6 / 16 / ncol / (nkptj + 1))) # +1 for pqkRbuf/pqkIbuf Gblksize = min(Gblksize, ngrids, 16384) pqkRbuf = numpy.empty(ncol * Gblksize) pqkIbuf = numpy.empty(ncol * Gblksize) buf = numpy.empty(nkptj * ncol * Gblksize, dtype=numpy.complex128) log.alldebug2('job_id %d blksize (%d,%d)', job_id, Gblksize, ncol) wcoulG = mydf.weighted_coulG(kpt, False, mesh) fused_cell_slice = (auxcell.nbas, fused_cell.nbas) if aosym == 's2': shls_slice = (sh0, sh1, 0, sh1) else: shls_slice = (sh0, sh1, 0, cell.nbas) for p0, p1 in lib.prange(0, ngrids, Gblksize): Gaux = ft_ao.ft_ao(fused_cell, Gv[p0:p1], fused_cell_slice, b, gxyz[p0:p1], Gvbase, kpt) Gaux *= wcoulG[p0:p1, None] kLR = Gaux.real.copy('C') kLI = Gaux.imag.copy('C') Gaux = None dat = ft_ao._ft_aopair_kpts(cell, Gv[p0:p1], shls_slice, aosym, b, gxyz[p0:p1], Gvbase, kpt, adapted_kptjs, out=buf) nG = p1 - p0 for k, ji in enumerate(adapted_ji_idx): aoao = dat[k].reshape(nG, ncol) pqkR = numpy.ndarray((ncol, nG), buffer=pqkRbuf) pqkI = numpy.ndarray((ncol, nG), buffer=pqkIbuf) pqkR[:] = aoao.real.T pqkI[:] = aoao.imag.T lib.dot(kLR.T, pqkR.T, -1, j3cR[k][naux:], 1) lib.dot(kLI.T, pqkI.T, -1, j3cR[k][naux:], 1) if not (is_zero(kpt) and gamma_point(adapted_kptjs[k])): lib.dot(kLR.T, pqkI.T, -1, j3cI[k][naux:], 1) lib.dot(kLI.T, pqkR.T, 1, j3cI[k][naux:], 1) kLR = kLI = None for k, idx in enumerate(adapted_ji_idx): if is_zero(kpt) and gamma_point(adapted_kptjs[k]): v = fuse(j3cR[k]) else: v = fuse(j3cR[k] + j3cI[k] * 1j) if j2ctag == 'CD': v = scipy.linalg.solve_triangular(j2c, v, lower=True, overwrite_b=True) fswap['j3c-chunks/%d/%d' % (job_id, idx)][:naux0] = v else: fswap['j3c-chunks/%d/%d' % (job_id, idx)][:naux0] = lib.dot( j2c, v) # low-dimension systems if j2c_negative is not None: fswap['j3c-/%d/%d' % (job_id, idx)] = lib.dot(j2c_negative, v)