def get_kconserv3(cell, kpts, kijkab): '''Get the momentum conservation array for a set of k-points. This function is similar to get_kconserv, but instead finds the 'kc' that satisfies momentum conservation for 5 k-points, kc = ki + kj + kk - ka - kb (mod G), where these kpoints are stored in kijkab[ki,kj,kk,ka,kb]. ''' nkpts = kpts.shape[0] KLMN = np.zeros([nkpts, nkpts, nkpts], np.int) kvecs = 2 * np.pi * scipy.linalg.inv(cell._h) kijkab = np.array(kijkab) # Finds which indices in ijkab are integers and which are lists # TODO: try to see if it works for more than 1 list idx_sum = np.array( [not (isinstance(x, int) or isinstance(x, np.int)) for x in kijkab]) idx_range = kijkab[idx_sum] min_idx_range = np.zeros(5, dtype=int) min_idx_range = np.array([min(x) for x in idx_range]) out_array_shape = tuple([len(x) for x in idx_range]) out_array = np.zeros(shape=out_array_shape, dtype=int) kpqrst_idx = np.zeros(5, dtype=int) # Order here matters! Search for most ``obvious" translation first to # get into 1st BZ, i.e. no translation! temp = [0, -1, 1, -2, 2] xyz = lib.cartesian_prod((temp, temp, temp)) kshift = np.dot(xyz, kvecs) for L, kvL in enumerate(lib.cartesian_prod(idx_range)): kpqrst_idx[idx_sum], kpqrst_idx[~idx_sum] = kvL, kijkab[~idx_sum] idx = tuple(kpqrst_idx[idx_sum] - min_idx_range) kvec = kpts[kpqrst_idx] kvec = kvec[0:3].sum(axis=0) - kvec[3:5].sum(axis=0) found = 0 kvNs = kvec + kshift for ishift in xrange(len(xyz)): kvN = kvNs[ishift] finder = np.where( np.logical_and(kpts < kvN + 1.e-12, kpts > kvN - 1.e-12).sum( axis=1) == 3) # The k-point kvN is the one that conserves momentum if len(finder[0]) > 0: found = 1 out_array[idx] = finder[0][0] break if found == 0: print "** ERROR: Problem in get_kconserv3. Quitting." print kijkab sys.exit() return out_array
def make_kpts(cell, nks, wrap_around=False, with_gamma_point=True): '''Given number of kpoints along x,y,z , generate kpoints Args: nks : (3,) ndarray Kwargs: wrap_around : bool To ensure all kpts are in first Brillouin zone. with_gamma_point : bool Whether to shift Monkhorst-pack grid to include gamma-point. Returns: kpts in absolute value (unit 1/Bohr). Gamma point is placed at the first place in the k-points list Examples: >>> cell.make_kpts((4,4,4)) ''' ks_each_axis = [] for n in nks: if with_gamma_point: ks = np.arange(n, dtype=float) / n else: ks = (np.arange(n) + .5) / n - .5 if wrap_around: ks[ks >= .5] -= 1 ks_each_axis.append(ks) scaled_kpts = lib.cartesian_prod(ks_each_axis) kpts = cell.get_abs_kpts(scaled_kpts) return kpts
def cell_plus_imgs(cell, nimgs): '''Create a supercell via nimgs[i] in each +/- direction, as in get_lattice_Ls(). Note this function differs from :fun:`super_cell` that super_cell only stacks the images in + direction. Args: cell : instance of :class:`Cell` nimgs : (3,) array Returns: supcell : instance of :class:`Cell` ''' supcell = cell.copy() a = cell.lattice_vectors() Ts = lib.cartesian_prod((np.arange(-nimgs[0],nimgs[0]+1), np.arange(-nimgs[1],nimgs[1]+1), np.arange(-nimgs[2],nimgs[2]+1))) Ls = np.dot(Ts, a) symbs = [atom[0] for atom in cell._atom] * len(Ls) coords = Ls.reshape(-1,1,3) + cell.atom_coords() supcell.atom = list(zip(symbs, coords.reshape(-1,3))) supcell.unit = 'B' supcell.a = np.einsum('i,ij->ij', nimgs, a) supcell.build(False, False, verbose=0) supcell.verbose = cell.verbose return supcell
def make_kpts(cell, nks, wrap_around=False): '''Given number of kpoints along x,y,z , generate kpoints Args: nks : (3,) ndarray Kwargs: wrap_around : bool To ensure all kpts are in first Brillouin zone. Returns: kpts in absolute value (unit 1/Bohr). Gamma point is placed at the first place in the k-points list Examples: >>> cell.make_kpts((4,4,4)) ''' ks_each_axis = [] for n in nks: ks = np.arange(n, dtype=float) / n if wrap_around: ks[ks>=.5] -= 1 ks_each_axis.append(ks) scaled_kpts = lib.cartesian_prod(ks_each_axis) kpts = cell.get_abs_kpts(scaled_kpts) return kpts
def ft_loop(self, mesh=None, q=numpy.zeros(3), kpts=None, shls_slice=None, max_memory=4000, aosym='s1', intor='GTO_ft_ovlp', comp=1, bvk_kmesh=None): ''' Fourier transform iterator for all kpti which satisfy 2pi*N = (kpts - kpti - q)*a, N = -1, 0, 1 ''' cell = self.cell if mesh is None: mesh = self.mesh if kpts is None: assert (is_zero(q)) kpts = self.kpts kpts = numpy.asarray(kpts) nkpts = len(kpts) ao_loc = cell.ao_loc_nr() b = cell.reciprocal_vectors() Gv, Gvbase, kws = cell.get_Gv_weights(mesh) gxyz = lib.cartesian_prod([numpy.arange(len(x)) for x in Gvbase]) ngrids = gxyz.shape[0] if shls_slice is None: shls_slice = (0, cell.nbas, 0, cell.nbas) if aosym == 's2': assert (shls_slice[2] == 0) i0 = ao_loc[shls_slice[0]] i1 = ao_loc[shls_slice[1]] nij = i1 * (i1 + 1) // 2 - i0 * (i0 + 1) // 2 else: 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 = max(16, int(max_memory * .9e6 / (nij * nkpts * 16 * comp))) blksize = min(blksize, ngrids, 16384) buf = numpy.empty(nkpts * nij * blksize * comp, dtype=numpy.complex128) for p0, p1 in self.prange(0, ngrids, blksize): dat = ft_ao.ft_aopair_kpts(cell, Gv[p0:p1], shls_slice, aosym, b, gxyz[p0:p1], Gvbase, q, kpts, intor, comp, bvk_kmesh=bvk_kmesh, out=buf) yield dat, p0, p1
def get_coords(self): """ Result: set of coordinates to compute a field which is to be stored in the file. """ coords = lib.cartesian_prod([self.xs, self.ys, self.zs]) coords = numpy.asarray(coords, order='C') - (-self.boxorig) return coords
def super_cell(cell, ncopy): '''Create an ncopy[0] x ncopy[1] x ncopy[2] supercell of the input cell Note this function differs from :fun:`cell_plus_imgs` that cell_plus_imgs creates images in both +/- direction. Args: cell : instance of :class:`Cell` ncopy : (3,) array Returns: supcell : instance of :class:`Cell` ''' a = cell.lattice_vectors() #:supcell.atom = [] #:for Lx in range(ncopy[0]): #: for Ly in range(ncopy[1]): #: for Lz in range(ncopy[2]): #: # Using cell._atom guarantees coord is in Bohr #: for atom, coord in cell._atom: #: L = np.dot([Lx, Ly, Lz], a) #: supcell.atom.append([atom, coord + L]) Ts = lib.cartesian_prod( (np.arange(ncopy[0]), np.arange(ncopy[1]), np.arange(ncopy[2]))) Ls = np.dot(Ts, a) supcell = cell.copy() supcell.a = np.einsum('i,ij->ij', ncopy, a) supcell.mesh = np.array([ ncopy[0] * cell.mesh[0], ncopy[1] * cell.mesh[1], ncopy[2] * cell.mesh[2] ]) return _build_supcell_(supcell, cell, Ls)
def get_coords(self): """ Result: set of coordinates to compute a field which is to be stored in the file. """ xyz = lib.cartesian_prod((self.xs, self.ys, self.zs)) coords = numpy.dot(xyz, self.box) return numpy.asarray(coords, order='C')
def get_lattice_Ls(cell, nimgs=None, rcut=None, dimension=None): '''Get the (Cartesian, unitful) lattice translation vectors for nearby images. The translation vectors can be used for the lattice summation.''' b = cell.reciprocal_vectors(norm_to=1) heights_inv = lib.norm(b, axis=1) if nimgs is None: if rcut is None: rcut = cell.rcut # plus 1 image in rcut to handle the case atoms within the adjacent cells are # close to each other rcut = rcut + min(1./heights_inv) nimgs = np.ceil(rcut*heights_inv) else: rcut = max((np.asarray(nimgs))/heights_inv) + min(1./heights_inv) # ~ the inradius if dimension is None: dimension = cell.dimension if dimension == 0: nimgs = [0, 0, 0] elif dimension == 1: nimgs = [nimgs[0], 0, 0] elif dimension == 2: nimgs = [nimgs[0], nimgs[1], 0] Ts = lib.cartesian_prod((np.arange(-nimgs[0],nimgs[0]+1), np.arange(-nimgs[1],nimgs[1]+1), np.arange(-nimgs[2],nimgs[2]+1))) Ls = np.dot(Ts, cell.lattice_vectors()) Ls = Ls[lib.norm(Ls, axis=1)<rcut] return np.asarray(Ls, order='C')
def cell_plus_imgs(cell, nimgs): '''Create a supercell via nimgs[i] in each +/- direction, as in get_lattice_Ls(). Note this function differs from :fun:`super_cell` that super_cell only stacks the images in + direction. Args: cell : instance of :class:`Cell` nimgs : (3,) array Returns: supcell : instance of :class:`Cell` ''' supcell = cell.copy() a = cell.lattice_vectors() Ts = lib.cartesian_prod( (np.arange(-nimgs[0], nimgs[0] + 1), np.arange(-nimgs[1], nimgs[1] + 1), np.arange(-nimgs[2], nimgs[2] + 1))) Ls = np.dot(Ts, a) symbs = [atom[0] for atom in cell._atom] * len(Ls) coords = Ls.reshape(-1, 1, 3) + cell.atom_coords() supcell.atom = list(zip(symbs, coords.reshape(-1, 3))) supcell.unit = 'B' supcell.a = np.einsum('i,ij->ij', nimgs, a) supcell.build(False, False, verbose=0) supcell.verbose = cell.verbose return supcell
def get_coords(self) : """ Result: set of coordinates to compute a field which is to be stored in the file. """ coords = lib.cartesian_prod([self.xs,self.ys,self.zs]) coords = numpy.asarray(coords, order='C') - (-self.boxorig) return coords
def _ao_rotation_matrices(mol, axes): '''Cache the rotation matrices''' from pyscf import lib from pyscf.symm.Dmatrix import Dmatrix, get_euler_angles alpha, beta, gamma = get_euler_angles(numpy.eye(3), axes) ANG_OF = 1 l_max = mol._bas[:,ANG_OF].max() if not mol.cart: return [Dmatrix(l, alpha, beta, gamma, reorder_p=True) for l in range(l_max+1)] pp = Dmatrix(1, alpha, beta, gamma, reorder_p=True) Ds = [numpy.ones((1,1))] for l in range(1, l_max+1): # All possible x,y,z combinations cidx = numpy.sort(lib.cartesian_prod([(0, 1, 2)] * l), axis=1) addr = 0 affine = numpy.ones((1,1)) for i in range(l): nd = affine.shape[0] * 3 affine = numpy.einsum('ik,jl->ijkl', affine, pp).reshape(nd, nd) addr = addr * 3 + cidx[:,i] uniq_addr, rev_addr = numpy.unique(addr, return_inverse=True) ncart = (l + 1) * (l + 2) // 2 assert ncart == uniq_addr.size trans = numpy.zeros((ncart,ncart)) for i, k in enumerate(rev_addr): trans[k] += affine[i,uniq_addr] Ds.append(trans) return Ds
def density(mol, outfile, dm, nx=80, ny=80, nz=80): coord = mol.atom_coords() box = numpy.max(coord, axis=0) - numpy.min(coord, axis=0) + 4 boxorig = numpy.min(coord, axis=0) - 2 xs = numpy.arange(nx) * (box[0] / nx) ys = numpy.arange(ny) * (box[1] / ny) zs = numpy.arange(nz) * (box[2] / nz) coords = lib.cartesian_prod([xs, ys, zs]) coords = numpy.asarray(coords, order="C") - (-boxorig) nao = mol.nao_nr() ngrids = nx * ny * nz blksize = min(200, ngrids) rho = numpy.empty(ngrids) for ip0, ip1 in gen_grid.prange(0, ngrids, blksize): ao = numint.eval_ao(mol, coords[ip0:ip1]) rho[ip0:ip1] = numint.eval_rho(mol, ao, dm) rho = rho.reshape(nx, ny, nz) with open(outfile, "w") as f: f.write("Density in real space\n") f.write("Comment line\n") f.write("%5d" % mol.natm) f.write(" %14.8f %14.8f %14.8f\n" % tuple(boxorig.tolist())) f.write("%5d %14.8f %14.8f %14.8f\n" % (nx, xs[1], 0, 0)) f.write("%5d %14.8f %14.8f %14.8f\n" % (ny, 0, ys[1], 0)) f.write("%5d %14.8f %14.8f %14.8f\n" % (nz, 0, 0, zs[1])) for ia in range(mol.natm): chg = mol.atom_charge(ia) f.write("%5d %f" % (chg, chg)) f.write(" %14.8f %14.8f %14.8f\n" % tuple(coord[ia])) fmt = " %14.8e" * nz + "\n" for ix in range(nx): for iy in range(ny): f.write(fmt % tuple(rho[ix, iy].tolist()))
def _discard_edge_images(cell, Ls, rcut): ''' Discard images if no basis in the image would contribute to lattice sum. ''' if rcut <= 0: return np.zeros((1, 3)) a = cell.lattice_vectors() scaled_atom_coords = np.linalg.solve(a.T, cell.atom_coords().T).T atom_boundary_max = scaled_atom_coords.max(axis=0) atom_boundary_min = scaled_atom_coords.min(axis=0) # ovlp_penalty ensures the overlap integrals for atoms in the adjcent # images are converged. ovlp_penalty = atom_boundary_max - atom_boundary_min # atom_boundary_min-1 ensures the values of basis at the grids on the edge # of the primitive cell converged boundary_max = np.ceil(np.max([atom_boundary_max, ovlp_penalty], axis=0)).astype(int) boundary_min = np.floor( np.min([atom_boundary_min - 1, -ovlp_penalty], axis=0)).astype(int) penalty_x = np.arange(boundary_min[0], boundary_max[0] + 1) penalty_y = np.arange(boundary_min[1], boundary_max[1] + 1) penalty_z = np.arange(boundary_min[2], boundary_max[2] + 1) shifts = lib.cartesian_prod([penalty_x, penalty_y, penalty_z]).dot(a) Ls_mask = (np.linalg.norm(Ls + shifts[:, None, :], axis=2) < rcut).any( axis=0) # cell0 (Ls == 0) should always be included. Ls_mask[len(Ls) // 2] = True return Ls[Ls_mask]
def get_phase(cell, kpts, kmesh=[]): ''' The unitary transformation that transforms the supercell basis k-mesh adapted basis. ''' latt_vec = cell.lattice_vectors() if kmesh is None: # Guess kmesh scaled_k = cell.get_scaled_kpts(kpts).round(8) kmesh = (len(numpy.unique(scaled_k[:,0])), len(numpy.unique(scaled_k[:,1])), len(numpy.unique(scaled_k[:,2]))) R_rel_a = numpy.arange(kmesh[0]) R_rel_b = numpy.arange(kmesh[1]) R_rel_c = numpy.arange(kmesh[2]) R_vec_rel = lib.cartesian_prod((R_rel_a, R_rel_b, R_rel_c)) R_vec_abs = numpy.einsum('nu, uv -> nv', R_vec_rel, latt_vec) NR = len(R_vec_abs) phase = numpy.exp(1j*numpy.einsum('Ru, ku -> Rk', R_vec_abs, kpts)) phase /= numpy.sqrt(NR) # normalization in supercell # R_rel_mesh has to be construct exactly same to the Ts in super_cell function scell = tools.super_cell(cell, kmesh) return scell, phase
def get_phase(cell, kpts, kmesh=None): ''' The unitary transformation that transforms the supercell basis k-mesh adapted basis. ''' latt_vec = cell.lattice_vectors() if kmesh is None: # Guess kmesh scaled_k = cell.get_scaled_kpts(kpts).round(8) kmesh = (len(np.unique(scaled_k[:,0])), len(np.unique(scaled_k[:,1])), len(np.unique(scaled_k[:,2]))) R_rel_a = np.arange(kmesh[0]) R_rel_b = np.arange(kmesh[1]) R_rel_c = np.arange(kmesh[2]) R_vec_rel = lib.cartesian_prod((R_rel_a, R_rel_b, R_rel_c)) R_vec_abs = np.einsum('nu, uv -> nv', R_vec_rel, latt_vec) NR = len(R_vec_abs) phase = np.exp(1j*np.einsum('Ru, ku -> Rk', R_vec_abs, kpts)) phase /= np.sqrt(NR) # normalization in supercell # R_rel_mesh has to be construct exactly same to the Ts in super_cell function scell = tools.super_cell(cell, kmesh) return scell, phase
def test_ft_aopair_bvk(self): from pyscf.pbc.tools import k2gamma n = 2 cell = pgto.Cell() cell.a = numpy.eye(3) * 4 cell.mesh = numpy.array([n, n, n]) cell.atom = '''C 1.3 .2 .3 C .1 .1 1.1 ''' cell.basis = 'ccpvdz' cell.unit = 'B' cell.build() kpts = cell.make_kpts([2, 2, 2]) Gv, Gvbase, kws = cell.get_Gv_weights() b = cell.reciprocal_vectors() gxyz = lib.cartesian_prod([numpy.arange(len(x)) for x in Gvbase]) bvk_kmesh = k2gamma.kpts_to_kmesh(cell, kpts) ref = ft_ao.ft_aopair_kpts(cell, Gv, b=b, gxyz=gxyz, Gvbase=Gvbase, kptjs=kpts) aopair = ft_ao.ft_aopair_kpts(cell, Gv, b=b, gxyz=gxyz, Gvbase=Gvbase, kptjs=kpts, bvk_kmesh=bvk_kmesh) self.assertAlmostEqual(abs(ref - aopair).max(), 0, 8) self.assertAlmostEqual(lib.fp(aopair), (-5.735639500461687 - 12.425151458809875j), 8)
def density_cut(mol, dm, nx=80, ny=80, z_value=0): from scipy.constants import physical_constants from pyscf import lib from pyscf.dft import gen_grid, numint nz = 1 coord = mol.atom_coords() box = np.max(coord,axis=0) - np.min(coord,axis=0) + 6 boxorig = np.min(coord,axis=0) - 3 xs = np.arange(nx) * (box[0]/nx) ys = np.arange(ny) * (box[1]/ny) zs = np.array([z_value]) #zs = np.arange(nz) * (box[2]/nz) coords = lib.cartesian_prod([xs,ys,zs]) coords = np.asarray(coords, order='C') - (-boxorig) ngrids = nx * ny * nz blksize = min(8000, ngrids) rho = np.empty(ngrids) ao = None for ip0, ip1 in gen_grid.prange(0, ngrids, blksize): ao = numint.eval_ao(mol, coords[ip0:ip1], out=ao) rho[ip0:ip1] = numint.eval_rho(mol, ao, dm) rho = rho.reshape(nx,ny) # needed for conversion as x,y are in bohr for some reason a0 = physical_constants["Bohr radius"][0] return rho.T, (xs + boxorig[0]) * a0, (ys + boxorig[1]) * a0
def get_lattice_Ls(cell, nimgs=None, rcut=None, dimension=None): '''Get the (Cartesian, unitful) lattice translation vectors for nearby images. The translation vectors can be used for the lattice summation.''' b = cell.reciprocal_vectors(norm_to=1) heights_inv = lib.norm(b, axis=1) if nimgs is None: if rcut is None: rcut = cell.rcut # plus 1 image in rcut to handle the case atoms within the adjacent cells are # close to each other rcut = rcut + max(1. / (heights_inv + 1e-8)) nimgs = np.ceil(rcut * heights_inv) else: rcut = max((np.asarray(nimgs)) / heights_inv) + min( 1. / heights_inv) # ~ the inradius if dimension is None: dimension = cell.dimension if dimension == 0: nimgs = [0, 0, 0] elif dimension == 1: nimgs = [nimgs[0], 0, 0] elif dimension == 2: nimgs = [nimgs[0], nimgs[1], 0] Ts = lib.cartesian_prod( (np.arange(-nimgs[0], nimgs[0] + 1), np.arange(-nimgs[1], nimgs[1] + 1), np.arange(-nimgs[2], nimgs[2] + 1))) Ls = np.dot(Ts, cell.lattice_vectors()) Ls = Ls[lib.norm(Ls, axis=1) < rcut] return np.asarray(Ls, order='C')
def precompute_exx(cell, kpts): from pyscf.pbc import gto as pbcgto from pyscf.pbc.dft import gen_grid log = lib.logger.Logger(cell.stdout, cell.verbose) log.debug("# Precomputing Wigner-Seitz EXX kernel") Nk = get_monkhorst_pack_size(cell, kpts) log.debug("# Nk = %s", Nk) kcell = pbcgto.Cell() kcell.atom = 'H 0. 0. 0.' kcell.spin = 1 kcell.unit = 'B' kcell.verbose = 0 kcell.a = cell.lattice_vectors() * Nk Lc = 1.0/lib.norm(np.linalg.inv(kcell.a), axis=0) log.debug("# Lc = %s", Lc) Rin = Lc.min() / 2.0 log.debug("# Rin = %s", Rin) # ASE: alpha = 5./Rin # sqrt(-ln eps) / Rc, eps ~ 10^{-11} log.info("WS alpha = %s", alpha) kcell.mesh = np.array([4*int(L*alpha*3.0) for L in Lc]) # ~ [120,120,120] # QE: #alpha = 3./Rin * np.sqrt(0.5) #kcell.mesh = (4*alpha*np.linalg.norm(kcell.a,axis=1)).astype(int) log.debug("# kcell.mesh FFT = %s", kcell.mesh) rs = gen_grid.gen_uniform_grids(kcell) kngs = len(rs) log.debug("# kcell kngs = %d", kngs) corners_coord = lib.cartesian_prod(([0, 1], [0, 1], [0, 1])) corners = np.dot(corners_coord, kcell.a) #vR = np.empty(kngs) #for i, rv in enumerate(rs): # # Minimum image convention to corners of kcell parallelepiped # r = lib.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 r = np.min([lib.norm(rs-c, axis=1) for c in corners], axis=0) vR = scipy.special.erf(alpha*r) / (r+1e-200) vR[r<1e-9] = 2*alpha / np.sqrt(np.pi) vG = (kcell.vol/kngs) * fft(vR, kcell.mesh) if abs(vG.imag).max() > 1e-6: # vG should be real in regular lattice. If imaginary part is observed, # this probably means a ws cell was built from a unconventional # lattice. The SR potential erfc(alpha*r) for the charge in the center # of ws cell decays to the region out of ws cell. The Ewald-sum based # on the minimum image convention cannot be used to build the kernel # Eq (12) of PRB 87, 165122 raise RuntimeError('Unconventional lattice was found') ws_exx = {'alpha': alpha, 'kcell': kcell, 'q' : kcell.Gv, 'vq' : vG.real.copy()} log.debug("# Finished precomputing") return ws_exx
def get_coords(self): xs = np.arange(self.nx) / (self.nx - 1) ys = np.arange(self.ny) / (self.ny - 1) zs = np.arange(self.nz) / (self.nz - 1) coords = lib.cartesian_prod([xs, ys, zs]) coords = np.dot(coords, self.a) coords = np.asarray(coords, order='C') + self.origin return coords
def setUpModule(): global mol, b, Gvbase, Gv, gxyz mol = gto.Mole() mol.atom = ''' C 1.3 .2 .3 C .1 -.1 1.1 ''' mol.basis = 'ccpvdz' mol.build() mesh = (7, 9, 11) numpy.random.seed(12) invh = numpy.diag(numpy.random.random(3)) b = 2 * numpy.pi * invh Gvbase = (numpy.fft.fftfreq(mesh[0], 1. / mesh[0]), numpy.fft.fftfreq(mesh[1], 1. / mesh[1]), numpy.fft.fftfreq(mesh[2], 1. / mesh[2])) Gv = numpy.dot(lib.cartesian_prod(Gvbase), b) gxyz = lib.cartesian_prod([numpy.arange(len(x)) for x in Gvbase])
def test_ft_aopair2(self): numpy.random.seed(12) invh = numpy.random.random(3) + numpy.eye(3) * 2.5 b = 2*numpy.pi * invh Gv = numpy.dot(lib.cartesian_prod(Gvbase), b) dat = ft_ao.ft_aopair(mol, Gv) self.assertAlmostEqual(finger(dat), (-3.1468496579780125-0.019209667673850885j), 9) dat1 = ft_ao.ft_aopair(mol, Gv, b=b, gxyz=gxyz, Gvbase=Gvbase) self.assertAlmostEqual(finger(dat1), (-3.1468496579780125-0.019209667673850885j), 9)
def get_coords(self): """ Result: set of coordinates to compute a field which is to be stored in the file. """ xs = numpy.arange(self.nx) * (1. / self.nx) ys = numpy.arange(self.ny) * (1. / self.ny) zs = numpy.arange(self.nz) * (1. / self.nz) xyz = lib.cartesian_prod((xs, ys, zs)) coords = numpy.dot(xyz, self.box) return numpy.asarray(coords, order='C')
def general_grid(cell, grid=[50, 50, 50]): ''' Generate a general grid for a cell object, this is different to the periodic PySCF grid object This general grid is used to generate the *.xsf file ''' ngrid = np.asarray(grid) qv = lib.cartesian_prod([np.arange(i) for i in ngrid]) a_frac = np.einsum('i,ij->ij', 1. / (ngrid - 1), cell.lattice_vectors()) coords = np.dot(qv, a_frac) return coords
def get_lattice_Ls(cell, nimgs=None): '''Get the (Cartesian, unitful) lattice translation vectors for nearby images.''' if nimgs is None: nimgs = cell.nimgs Ts = lib.cartesian_prod((np.arange(-nimgs[0],nimgs[0]+1), np.arange(-nimgs[1],nimgs[1]+1), np.arange(-nimgs[2],nimgs[2]+1))) #Ts = Ts[np.einsum('ix,ix->i',Ts,Ts) <= 1./3*np.dot(nimgs,nimgs)] Ts = Ts[np.einsum('ix,ix->i',Ts,Ts) <= max(nimgs)*max(nimgs)] Ls = np.dot(Ts, cell._h.astype(np.double).T) return Ls
def translation_vectors_for_kmesh(cell, kmesh): ''' Translation vectors to construct super-cell of which the gamma point is identical to the k-point mesh of primitive cell ''' latt_vec = cell.lattice_vectors() R_rel_a = np.arange(kmesh[0]) R_rel_b = np.arange(kmesh[1]) R_rel_c = np.arange(kmesh[2]) R_vec_rel = lib.cartesian_prod((R_rel_a, R_rel_b, R_rel_c)) R_vec_abs = np.einsum('nu, uv -> nv', R_vec_rel, latt_vec) return R_vec_abs
def get_Gv_weights(cell, gs=None): '''Calculate G-vectors and weights. Returns: Gv : (ngs, 3) ndarray of floats The array of G-vectors. ''' if gs is None: gs = cell.gs def plus_minus(n): #rs, ws = dft.delley(n) #rs, ws = dft.treutler_ahlrichs(n) #rs, ws = dft.mura_knowles(n) rs, ws = dft.gauss_chebyshev(n) #return np.hstack((0,rs,-rs[::-1])), np.hstack((0,ws,ws[::-1])) return np.hstack((rs, -rs[::-1])), np.hstack((ws, ws[::-1])) # Default, the 3D uniform grids b = cell.reciprocal_vectors() rx = np.append(np.arange(gs[0] + 1.), np.arange(-gs[0], 0.)) ry = np.append(np.arange(gs[1] + 1.), np.arange(-gs[1], 0.)) rz = np.append(np.arange(gs[2] + 1.), np.arange(-gs[2], 0.)) ngs = [i * 2 + 1 for i in gs] if cell.dimension == 0: rx, wx = plus_minus(gs[0]) ry, wy = plus_minus(gs[1]) rz, wz = plus_minus(gs[2]) rx /= np.linalg.norm(b[0]) ry /= np.linalg.norm(b[1]) rz /= np.linalg.norm(b[2]) weights = np.einsum('i,j,k->ijk', wx, wy, wz).reshape(-1) elif cell.dimension == 1: wx = np.repeat(np.linalg.norm(b[0]), ngs[0]) ry, wy = plus_minus(gs[1]) rz, wz = plus_minus(gs[2]) ry /= np.linalg.norm(b[1]) rz /= np.linalg.norm(b[2]) weights = np.einsum('i,j,k->ijk', wx, wy, wz).reshape(-1) elif cell.dimension == 2: area = np.linalg.norm(np.cross(b[0], b[1])) wxy = np.repeat(area, ngs[0] * ngs[1]) rz, wz = plus_minus(gs[2]) rz /= np.linalg.norm(b[2]) weights = np.einsum('i,k->ik', wxy, wz).reshape(-1) else: weights = abs(np.linalg.det(b)) Gvbase = (rx, ry, rz) Gv = np.dot(lib.cartesian_prod(Gvbase), b) # 1/cell.vol == det(b)/(2pi)^3 weights *= 1 / (2 * np.pi)**3 return Gv, Gvbase, weights
def get_lattice_Ls(cell, nimgs=None): '''Get the (Cartesian, unitful) lattice translation vectors for nearby images.''' if nimgs is None: nimgs = cell.nimgs Ts = lib.cartesian_prod( (np.arange(-nimgs[0], nimgs[0] + 1), np.arange(-nimgs[1], nimgs[1] + 1), np.arange(-nimgs[2], nimgs[2] + 1))) #Ts = Ts[np.einsum('ix,ix->i',Ts,Ts) <= 1./3*np.dot(nimgs,nimgs)] Ts = Ts[np.einsum('ix,ix->i', Ts, Ts) <= max(nimgs) * max(nimgs)] Ls = np.dot(Ts, cell._h.astype(np.double).T) return Ls
def get_Gv_weights(cell, gs=None): '''Calculate G-vectors and weights. Returns: Gv : (ngs, 3) ndarray of floats The array of G-vectors. ''' if gs is None: gs = cell.gs def plus_minus(n): #rs, ws = dft.delley(n) #rs, ws = dft.treutler_ahlrichs(n) #rs, ws = dft.mura_knowles(n) rs, ws = dft.gauss_chebyshev(n) #return np.hstack((0,rs,-rs[::-1])), np.hstack((0,ws,ws[::-1])) return np.hstack((rs,-rs[::-1])), np.hstack((ws,ws[::-1])) # Default, the 3D uniform grids b = cell.reciprocal_vectors() rx = np.append(np.arange(gs[0]+1.), np.arange(-gs[0],0.)) ry = np.append(np.arange(gs[1]+1.), np.arange(-gs[1],0.)) rz = np.append(np.arange(gs[2]+1.), np.arange(-gs[2],0.)) weights = np.linalg.det(b) ngs = [i*2+1 for i in gs] if cell.dimension == 0: rx, wx = plus_minus(gs[0]) ry, wy = plus_minus(gs[1]) rz, wz = plus_minus(gs[2]) rx /= np.linalg.norm(b[0]) ry /= np.linalg.norm(b[1]) rz /= np.linalg.norm(b[2]) weights = np.einsum('i,j,k->ijk', wx, wy, wz).reshape(-1) elif cell.dimension == 1: wx = np.repeat(np.linalg.norm(b[0]), ngs[0]) ry, wy = plus_minus(gs[1]) rz, wz = plus_minus(gs[2]) ry /= np.linalg.norm(b[1]) rz /= np.linalg.norm(b[2]) weights = np.einsum('i,j,k->ijk', wx, wy, wz).reshape(-1) elif cell.dimension == 2: area = np.linalg.norm(np.cross(b[0], b[1])) wxy = np.repeat(area, ngs[0]*ngs[1]) rz, wz = plus_minus(gs[2]) rz /= np.linalg.norm(b[2]) weights = np.einsum('i,k->ik', wxy, wz).reshape(-1) Gvbase = (rx, ry, rz) Gv = np.dot(lib.cartesian_prod(Gvbase), b) # 1/cell.vol == det(b)/(2pi)^3 weights *= 1/(2*np.pi)**3 return Gv, Gvbase, weights
def test_ft_aopair2(self): numpy.random.seed(12) invh = numpy.random.random(3) + numpy.eye(3) * 2.5 b = 2 * numpy.pi * invh Gv = numpy.dot(lib.cartesian_prod(Gvbase), b) dat = ft_ao.ft_aopair(mol, Gv) self.assertAlmostEqual(finger(dat), (-3.1468496579780125 - 0.019209667673850885j), 9) dat1 = ft_ao.ft_aopair(mol, Gv, b=b, gxyz=gxyz, Gvbase=Gvbase) self.assertAlmostEqual(finger(dat1), (-3.1468496579780125 - 0.019209667673850885j), 9)
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 test_ft_ao2(self): numpy.random.seed(12) invh = numpy.random.random(3) + numpy.eye(3) * 2.5 b = 2 * numpy.pi * invh Gv = numpy.dot(lib.cartesian_prod(Gvbase), b) ref = ft_ao_o0(mol, Gv) dat = ft_ao.ft_ao(mol, Gv) self.assertTrue(numpy.allclose(ref, dat)) mol1 = mol.copy() mol1.cart = True ref = ft_ao.ft_ao(mol1, Gv) dat = ft_ao.ft_ao(mol1, Gv, b=b, gxyz=gxyz, Gvbase=Gvbase) self.assertTrue(numpy.allclose(ref, dat))
def test_ft_ao2(self): numpy.random.seed(12) invh = numpy.random.random(3) + numpy.eye(3) * 2.5 b = 2*numpy.pi * invh Gv = numpy.dot(lib.cartesian_prod(Gvbase), b) ref = ft_ao_o0(mol, Gv) dat = ft_ao.ft_ao(mol, Gv) self.assertTrue(numpy.allclose(ref, dat)) mol1 = mol.copy() mol1.cart = True ref = ft_ao.ft_ao(mol1, Gv) dat = ft_ao.ft_ao(mol1, Gv, b=b, gxyz=gxyz, Gvbase=Gvbase) self.assertTrue(numpy.allclose(ref, dat))
def test_ft_aoao(self): #coords = pdft.gen_grid.gen_uniform_grids(cell) #aoR = pdft.numint.eval_ao(cell, coords) #ngs, nao = aoR.shape #ref = numpy.asarray([tools.fft(aoR[:,i].conj()*aoR[:,j], cell.gs) # for i in range(nao) for j in range(nao)]) #ref = ref.reshape(nao,nao,-1).transpose(2,0,1) * (cell.vol/ngs) #dat = ft_ao.ft_aopair(cell, cell.Gv, aosym='s1hermi') #self.assertAlmostEqual(numpy.linalg.norm(ref[:,0,0]-dat[:,0,0]) , 0, 5) #self.assertAlmostEqual(numpy.linalg.norm(ref[:,1,1]-dat[:,1,1]) , 0.02315483195832373, 4) #self.assertAlmostEqual(numpy.linalg.norm(ref[:,2:,2:]-dat[:,2:,2:]), 0, 9) #self.assertAlmostEqual(numpy.linalg.norm(ref[:,0,2:]-dat[:,0,2:]) , 0, 9) #self.assertAlmostEqual(numpy.linalg.norm(ref[:,2:,0]-dat[:,2:,0]) , 0, 9) #idx = numpy.tril_indices(nao) #ref = dat[:,idx[0],idx[1]] #dat = ft_ao.ft_aopair(cell, cell.Gv, aosym='s2') #self.assertAlmostEqual(abs(dat-ref).sum(), 0, 9) coords = pdft.gen_grid.gen_uniform_grids(cell1) aoR = pdft.numint.eval_ao(cell1, coords) ngs, nao = aoR.shape ref = numpy.asarray([ tools.fft(aoR[:, i].conj() * aoR[:, j], cell1.gs) for i in range(nao) for j in range(nao) ]) ref = ref.reshape(nao, nao, -1).transpose(2, 0, 1) * (cell1.vol / ngs) Gv, Gvbase, kws = cell1.get_Gv_weights(cell1.gs) b = cell1.reciprocal_vectors() gxyz = lib.cartesian_prod([numpy.arange(len(x)) for x in Gvbase]) dat = ft_ao.ft_aopair(cell1, cell1.Gv, aosym='s1hermi', b=b, gxyz=gxyz, Gvbase=Gvbase) self.assertAlmostEqual(numpy.linalg.norm(ref[:, 0, 0] - dat[:, 0, 0]), 0, 7) self.assertAlmostEqual(numpy.linalg.norm(ref[:, 1, 1] - dat[:, 1, 1]), 0, 7) self.assertAlmostEqual( numpy.linalg.norm(ref[:, 2:, 2:] - dat[:, 2:, 2:]), 0, 7) self.assertAlmostEqual( numpy.linalg.norm(ref[:, 0, 2:] - dat[:, 0, 2:]), 0, 7) self.assertAlmostEqual( numpy.linalg.norm(ref[:, 2:, 0] - dat[:, 2:, 0]), 0, 7) idx = numpy.tril_indices(nao) ref = dat[:, idx[0], idx[1]] dat = ft_ao.ft_aopair(cell1, cell1.Gv, aosym='s2') self.assertAlmostEqual(abs(dat - ref).sum(), 0, 9)
def non_uniform_kgrids(gs): from pyscf.dft import gen_grid def plus_minus(n): #rs, ws = gen_grid.radi.delley(n) #rs, ws = gen_grid.radi.treutler_ahlrichs(n) #rs, ws = gen_grid.radi.mura_knowles(n) rs, ws = gen_grid.radi.gauss_chebyshev(n) return numpy.hstack((rs,-rs)), numpy.hstack((ws,ws)) rx, wx = plus_minus(gs[0]) ry, wy = plus_minus(gs[1]) rz, wz = plus_minus(gs[2]) Gvbase = (rx, ry, rz) Gv = lib.cartesian_prod(Gvbase) weights = numpy.einsum('i,j,k->ijk', wx, wy, wz).reshape(-1) return Gv, Gvbase, weights
def get_wannier(w90, supercell=[1, 1, 1], grid=[50, 50, 50]): ''' Evaluate the MLWF using a periodic grid ''' import sys sys.path.append('/home/gagliard/phamx494/CPPlib/pyWannier90') import libwannier90 sys.path.append('/panfs/roc/groups/6/gagliard/phamx494/CPPlib/pyWannier90') import pywannier90 from pyscf.pbc.dft import gen_grid, numint grids_coor, weights = pywannier90.periodic_grid(w90.cell, grid, supercell=[1, 1, 1], order='C') kpts = w90.cell.get_abs_kpts(w90.kpt_latt_loc) ao_kpts = np.asarray( [numint.eval_ao(w90.cell, grids_coor, kpt=kpt) for kpt in kpts]) u_mo = [] for k_id in range(w90.num_kpts_loc): mo_included = w90.mo_coeff_kpts[k_id][:, w90.band_included_list] mo_in_window = w90.lwindow[k_id] C_opt = mo_included[:, mo_in_window].dot(w90.U_matrix_opt[k_id].T) C_tildle = C_opt.dot(w90.U_matrix[k_id].T) kpt = kpts[k_id] ao = numint.eval_ao(w90.cell, grids_coor, kpt=kpt) u_ao = np.einsum('x,xi->xi', np.exp(-1j * np.dot(grids_coor, kpt)), ao, optimize=True) u_mo.append(np.einsum('xi,in->xn', u_ao, C_tildle, optimize=True)) u_mo = np.asarray(u_mo) nimgs = [kpt // 2 for kpt in w90.mp_grid_loc] Ts = lib.cartesian_prod( (np.arange(-nimgs[0], nimgs[0] + 1), np.arange(-nimgs[1], nimgs[1] + 1), np.arange(-nimgs[2], nimgs[2] + 1))) Ts = np.asarray( Ts, order='C') #lib.cartesian_prod store array in Fortran order in memory WFs = libwannier90.get_WFs(w90.kpt_latt_loc.shape[0], w90.kpt_latt_loc, Ts.shape[0], Ts, supercell, grid, u_mo) return WFs
def make_kpts(cell, nks): '''Given number of kpoints along x,y,z , generate kpoints Args: nks : (3,) ndarray Returns: kpts in absolute value (unit 1/Bohr) Examples: >>> cell.make_kpts((4,4,4)) ''' ks_each_axis = [(np.arange(n)+.5)/n-.5 for n in nks] scaled_kpts = lib.cartesian_prod(ks_each_axis) kpts = cell.get_abs_kpts(scaled_kpts) return kpts
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 ft_loop(self, cell, gs=None, kpt=numpy.zeros(3), kpts=None, shls_slice=None, max_memory=4000): ''' Fourier transform iterator for all kpti which satisfy kpt = kpts - kpti ''' if gs is None: gs = self.gs if kpts is None: assert(gamma_point(kpt)) kpts = self.kpts kpts = numpy.asarray(kpts) nkpts = len(kpts) 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(kpt) and gamma_point(kpts): aosym = 's1hermi' else: aosym = 's1' blksize = min(max(16, int(max_memory*.9e6/(nao**2*(nkpts+1)*16))), 16384) buf = [numpy.zeros(nao*nao*blksize, dtype=numpy.complex128) for k in range(nkpts)] pqkRbuf = numpy.empty(nao*nao*blksize) pqkIbuf = numpy.empty(nao*nao*blksize) for p0, p1 in self.prange(0, ngs, blksize): ft_ao._ft_aopair_kpts(cell, Gv[p0:p1], shls_slice, aosym, invh, gxyz[p0:p1], gs, kpt, kpts, out=buf) nG = p1 - p0 for k in range(nkpts): aoao = numpy.ndarray((nG,nao,nao), dtype=numpy.complex128, order='F', buffer=buf[k]) pqkR = numpy.ndarray((nao,nao,nG), buffer=pqkRbuf) pqkI = numpy.ndarray((nao,nao,nG), buffer=pqkIbuf) pqkR[:] = aoao.real.transpose(1,2,0) pqkI[:] = aoao.imag.transpose(1,2,0) yield (k, pqkR.reshape(-1,nG), pqkI.reshape(-1,nG), p0, p1) aoao[:] = 0 # == buf[k][:] = 0
def get_phase(self, cell=None, kpts=None, kmesh=None): ''' Get a super cell and the phase matrix that transform from real to k-space ''' if kmesh is None: kmesh = w90.mp_grid_loc if cell is None: cell = self.cell if kpts is None: kpts = self.kpts a = cell.lattice_vectors() Ts = lib.cartesian_prod( (np.arange(kmesh[0]), np.arange(kmesh[1]), np.arange(kmesh[2]))) Rs = np.dot(Ts, a) NRs = Rs.shape[0] phase = 1 / np.sqrt(NRs) * np.exp(1j * Rs.dot(kpts.T)) scell = pbctools.super_cell(cell, kmesh) return scell, phase
def gen_uniform_grids(cell, gs=None): '''Generate a uniform real-space grid consistent w/ samp thm; see MH (3.19). Args: cell : instance of :class:`Cell` Returns: coords : (ngx*ngy*ngz, 3) ndarray The real-space grid point coordinates. ''' if gs is None: gs = cell.gs ngs = 2*np.asarray(gs)+1 qv = lib.cartesian_prod([np.arange(x) for x in ngs]) a_frac = np.einsum('i,ij->ij', 1./ngs, cell.lattice_vectors()) coords = np.dot(qv, a_frac) return coords
def gen_uniform_grids(cell, gs=None): '''Generate a uniform real-space grid consistent w/ samp thm; see MH (3.19). Args: cell : instance of :class:`Cell` Returns: coords : (ngx*ngy*ngz, 3) ndarray The real-space grid point coordinates. ''' if gs is None: gs = cell.gs ngs = 2 * np.asarray(gs) + 1 qv = lib.cartesian_prod([np.arange(x) for x in ngs]) a_frac = np.einsum('i,ij->ij', 1. / ngs, cell.lattice_vectors()) coords = np.dot(qv, a_frac) return coords
def make_kpts(cell, nks): '''Given number of kpoints along x,y,z , generate kpoints Args: nks : (3,) ndarray Returns: kpts in absolute value (unit 1/Bohr). Gamma point is placed at the first place in the k-points list Examples: >>> cell.make_kpts((4,4,4)) ''' ks_each_axis = [np.arange(n,dtype=float)/n for n in nks] scaled_kpts = lib.cartesian_prod(ks_each_axis) kpts = cell.get_abs_kpts(scaled_kpts) return kpts
def __init__(self, cell, kpts): '''Helper class for handling k-points in correlated calculations. Attributes: kconserv : (nkpts,nkpts,nkpts) ndarray The index of the fourth momentum-conserving k-point, given indices of three k-points symm_map : OrderedDict of list of (3,) tuples Keys are (3,) tuples of symmetry-unique k-point indices and values are lists of (3,) tuples, enumerating all symmetry-related k-point indices for ERI generation ''' self.kconserv = get_kconserv(cell, kpts) nkpts = len(kpts) temp = range(0,nkpts) kptlist = lib.cartesian_prod((temp,temp,temp)) completed = np.zeros((nkpts,nkpts,nkpts), dtype=bool) self._operation = np.zeros((nkpts,nkpts,nkpts), dtype=int) self.symm_map = OrderedDict() for kpt in kptlist: kpt = tuple(kpt) kp,kq,kr = kpt if not completed[kp,kq,kr]: self.symm_map[kpt] = list() ks = self.kconserv[kp,kq,kr] completed[kp,kq,kr] = True self._operation[kp,kq,kr] = 0 self.symm_map[kpt].append((kp,kq,kr)) completed[kr,ks,kp] = True self._operation[kr,ks,kp] = 1 #.transpose(2,3,0,1) self.symm_map[kpt].append((kr,ks,kp)) completed[kq,kp,ks] = True self._operation[kq,kp,ks] = 2 #np.conj(.transpose(1,0,3,2)) self.symm_map[kpt].append((kq,kp,ks)) completed[ks,kr,kq] = True self._operation[ks,kr,kq] = 3 #np.conj(.transpose(3,2,1,0)) self.symm_map[kpt].append((ks,kr,kq))
def ft_loop(self, mesh=None, q=numpy.zeros(3), kpts=None, shls_slice=None, max_memory=4000, aosym='s1', intor='GTO_ft_ovlp', comp=1): ''' Fourier transform iterator for all kpti which satisfy 2pi*N = (kpts - kpti - q)*a, N = -1, 0, 1 ''' cell = self.cell if mesh is None: mesh = self.mesh if kpts is None: assert(is_zero(q)) kpts = self.kpts kpts = numpy.asarray(kpts) nkpts = len(kpts) ao_loc = cell.ao_loc_nr() b = cell.reciprocal_vectors() Gv, Gvbase, kws = cell.get_Gv_weights(mesh) gxyz = lib.cartesian_prod([numpy.arange(len(x)) for x in Gvbase]) ngrids = gxyz.shape[0] if shls_slice is None: shls_slice = (0, cell.nbas, 0, cell.nbas) if aosym == 's2': assert(shls_slice[2] == 0) i0 = ao_loc[shls_slice[0]] i1 = ao_loc[shls_slice[1]] nij = i1*(i1+1)//2 - i0*(i0+1)//2 else: 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 = max(16, int(max_memory*.9e6/(nij*nkpts*16*comp))) blksize = min(blksize, ngrids, 16384) buf = numpy.empty(nkpts*nij*blksize*comp, dtype=numpy.complex128) for p0, p1 in self.prange(0, ngrids, blksize): dat = ft_ao._ft_aopair_kpts(cell, Gv[p0:p1], shls_slice, aosym, b, gxyz[p0:p1], Gvbase, q, kpts, intor, comp, out=buf) yield dat, p0, p1
def get_Gv(self): """Calculate three-dimensional G-vectors for the cell; see MH (3.8). Indices along each direction go as [0...self.gs, -self.gs...-1] to follow FFT convention. Note that, for each direction, ngs = 2*self.gs+1. Args: self : instance of :class:`Cell` Returns: Gv : (ngs, 3) ndarray of floats The array of G-vectors. """ gxrange = range(self.gs[0] + 1) + range(-self.gs[0], 0) gyrange = range(self.gs[1] + 1) + range(-self.gs[1], 0) gzrange = range(self.gs[2] + 1) + range(-self.gs[2], 0) gxyz = lib.cartesian_prod((gxrange, gyrange, gzrange)) invh = scipy.linalg.inv(self.lattice_vectors()) Gv = 2 * np.pi * np.dot(gxyz, invh) return Gv
def get_kconserv(cell, kpts): '''Get the momentum conservation array for a set of k-points. Given k-point indices (k, l, m) the array kconserv[k,l,m] returns the index n that satifies momentum conservation, k(k) - k(l) = - k(m) + k(n) This is used for symmetry e.g. integrals of the form [\phi*[k](1) \phi[l](1) | \phi*[m](2) \phi[n](2)] are zero unless n satisfies the above. ''' nkpts = kpts.shape[0] KLMN = np.zeros([nkpts,nkpts,nkpts], np.int) kvecs = cell.reciprocal_vectors() for K, kvK in enumerate(kpts): for L, kvL in enumerate(kpts): for M, kvM in enumerate(kpts): # Here we find where kvN = kvM + kvL - kvK (mod K) temp = range(-1,2) xyz = lib.cartesian_prod((temp,temp,temp)) found = 0 kvMLK = kvK - kvL + kvM kvN = kvMLK for ishift in xrange(len(xyz)): kvN = kvMLK + np.dot(xyz[ishift],kvecs) finder = np.where(np.logical_and(kpts < kvN + 1.e-12, kpts > kvN - 1.e-12).sum(axis=1)==3) # The k-point should be the same in all 3 indices as kvN if len(finder[0]) > 0: KLMN[K, L, M] = finder[0][0] found = 1 break if found == 0: print("** ERROR: Problem in get_kconserv. Quitting.") print(kvMLK) sys.exit() return KLMN
def get_Gv(cell, gs=None): '''Calculate three-dimensional G-vectors for the cell; see MH (3.8). Indices along each direction go as [0...cell.gs, -cell.gs...-1] to follow FFT convention. Note that, for each direction, ngs = 2*cell.gs+1. Args: cell : instance of :class:`Cell` Returns: Gv : (ngs, 3) ndarray of floats The array of G-vectors. ''' if gs is None: gs = cell.gs gxrange = np.append(range(gs[0]+1), range(-gs[0],0)) gyrange = np.append(range(gs[1]+1), range(-gs[1],0)) gzrange = np.append(range(gs[2]+1), range(-gs[2],0)) gxyz = lib.cartesian_prod((gxrange, gyrange, gzrange)) invh = scipy.linalg.inv(cell.lattice_vectors()) Gv = 2*np.pi* np.dot(gxyz, invh) return Gv
def super_cell(cell, ncopy): '''Create an ncopy[0] x ncopy[1] x ncopy[2] supercell of the input cell Note this function differs from :fun:`cell_plus_imgs` that cell_plus_imgs creates images in both +/- direction. Args: cell : instance of :class:`Cell` ncopy : (3,) array Returns: supcell : instance of :class:`Cell` ''' supcell = cell.copy() a = cell.lattice_vectors() #:supcell.atom = [] #:for Lx in range(ncopy[0]): #: for Ly in range(ncopy[1]): #: for Lz in range(ncopy[2]): #: # Using cell._atom guarantees coord is in Bohr #: for atom, coord in cell._atom: #: L = np.dot([Lx, Ly, Lz], a) #: supcell.atom.append([atom, coord + L]) Ts = lib.cartesian_prod((np.arange(ncopy[0]), np.arange(ncopy[1]), np.arange(ncopy[2]))) Ls = np.dot(Ts, a) symbs = [atom[0] for atom in cell._atom] * len(Ls) coords = Ls.reshape(-1,1,3) + cell.atom_coords() supcell.atom = list(zip(symbs, coords.reshape(-1,3))) supcell.unit = 'B' supcell.a = np.einsum('i,ij->ij', ncopy, a) supcell.mesh = np.array([ncopy[0]*cell.mesh[0], ncopy[1]*cell.mesh[1], ncopy[2]*cell.mesh[2]]) supcell.build(False, False, verbose=0) supcell.verbose = cell.verbose return supcell
def get_lattice_Ls(cell, nimgs=None, rcut=None, dimension=None): '''Get the (Cartesian, unitful) lattice translation vectors for nearby images. The translation vectors can be used for the lattice summation.''' a = cell.lattice_vectors() b = cell.reciprocal_vectors(norm_to=1) heights_inv = lib.norm(b, axis=1) if nimgs is None: if rcut is None: rcut = cell.rcut # plus 1 image in rcut to handle the case atoms within the adjacent cells are # close to each other nimgs = np.ceil(rcut*heights_inv + 1.1).astype(int) else: rcut = max((np.asarray(nimgs))/heights_inv) if dimension is None: dimension = cell.dimension if dimension == 0: nimgs = [0, 0, 0] elif dimension == 1: nimgs = [nimgs[0], 0, 0] elif dimension == 2: nimgs = [nimgs[0], nimgs[1], 0] Ts = lib.cartesian_prod((np.arange(-nimgs[0],nimgs[0]+1), np.arange(-nimgs[1],nimgs[1]+1), np.arange(-nimgs[2],nimgs[2]+1))) Ls = np.dot(Ts, a) idx = np.zeros(len(Ls), dtype=bool) for ax in (-a[0], 0, a[0]): for ay in (-a[1], 0, a[1]): for az in (-a[2], 0, a[2]): idx |= lib.norm(Ls+(ax+ay+az), axis=1) < rcut Ls = Ls[idx] return np.asarray(Ls, order='C')
def _make_j3c(mydf, cell, auxcell, kptij_lst, cderi_file): log = logger.Logger(mydf.stdout, mydf.verbose) t1 = t0 = (time.clock(), time.time()) fused_cell, fuse = fuse_auxcell(mydf, mydf.auxcell) ao_loc = cell.ao_loc_nr() nao = ao_loc[-1] naux = auxcell.nao_nr() nkptij = len(kptij_lst) mesh = mydf.mesh Gv, Gvbase, kws = cell.get_Gv_weights(mesh) b = cell.reciprocal_vectors() gxyz = lib.cartesian_prod([numpy.arange(len(x)) for x in Gvbase]) ngrids = gxyz.shape[0] kptis = kptij_lst[:,0] kptjs = kptij_lst[:,1] kpt_ji = kptjs - kptis uniq_kpts, uniq_index, uniq_inverse = unique(kpt_ji) log.debug('Num uniq kpts %d', len(uniq_kpts)) log.debug2('uniq_kpts %s', uniq_kpts) # j2c ~ (-kpt_ji | kpt_ji) j2c = fused_cell.pbc_intor('int2c2e', hermi=1, kpts=uniq_kpts) j2ctags = [] t1 = log.timer_debug1('2c2e', *t1) swapfile = tempfile.NamedTemporaryFile(dir=os.path.dirname(cderi_file)) fswap = lib.H5TmpFile(swapfile.name) # Unlink swapfile to avoid trash swapfile = None for k, kpt in enumerate(uniq_kpts): coulG = mydf.weighted_coulG(kpt, False, mesh) j2c[k] = fuse(fuse(j2c[k]).T).T.copy() j2c_k = numpy.zeros_like(j2c[k]) for p0, p1 in mydf.mpi_prange(0, ngrids): aoaux = ft_ao.ft_ao(fused_cell, Gv[p0:p1], None, b, gxyz[p0:p1], Gvbase, kpt).T aoaux = fuse(aoaux) LkR = numpy.asarray(aoaux.real, order='C') LkI = numpy.asarray(aoaux.imag, order='C') aoaux = None if is_zero(kpt): # kpti == kptj j2cR = lib.dot(LkR*coulG[p0:p1], LkR.T) j2c_k += lib.dot(LkI*coulG[p0:p1], LkI.T, 1, j2cR, 1) else: # aoaux ~ kpt_ij, aoaux.conj() ~ kpt_kl j2cR, j2cI = zdotCN(LkR*coulG[p0:p1], LkI*coulG[p0:p1], LkR.T, LkI.T) j2c_k += j2cR + j2cI * 1j LkR = LkI = None j2c[k] -= mpi.allreduce(j2c_k) try: fswap['j2c/%d'%k] = scipy.linalg.cholesky(j2c[k], lower=True) j2ctags.append('CD') except scipy.linalg.LinAlgError: w, v = scipy.linalg.eigh(j2c[k]) log.debug2('metric linear dependency for kpt %s', k) log.debug2('cond = %.4g, drop %d bfns', w[0]/w[-1], numpy.count_nonzero(w<mydf.linear_dep_threshold)) v1 = v[:,w>mydf.linear_dep_threshold].T.conj() v1 /= numpy.sqrt(w[w>mydf.linear_dep_threshold]).reshape(-1,1) fswap['j2c/%d'%k] = v1 if cell.dimension == 2 and cell.low_dim_ft_type != 'inf_vacuum': idx = numpy.where(w < -mydf.linear_dep_threshold)[0] if len(idx) > 0: fswap['j2c-/%d'%k] = (v[:,idx]/numpy.sqrt(-w[idx])).conj().T w = v = v1 = v2 = None j2ctags.append('eig') aoaux = kLR = kLI = j2cR = j2cI = coulG = None j2c = None aosym_s2 = numpy.einsum('ix->i', abs(kptis-kptjs)) < 1e-9 j_only = numpy.all(aosym_s2) if gamma_point(kptij_lst): dtype = 'f8' else: dtype = 'c16' t1 = log.timer_debug1('aoaux and int2c', *t1) # Estimates the buffer size based on the last contraction in G-space. # This contraction requires to hold nkptj copies of (naux,?) array # simultaneously in memory. mem_now = max(comm.allgather(lib.current_memory()[0])) max_memory = max(2000, mydf.max_memory - mem_now) nkptj_max = max((uniq_inverse==x).sum() for x in set(uniq_inverse)) buflen = max(int(min(max_memory*.5e6/16/naux/(nkptj_max+2)/nao, nao/3/mpi.pool.size)), 1) chunks = (buflen, nao) j3c_jobs = mpi_df.grids2d_int3c_jobs(cell, auxcell, kptij_lst, chunks, j_only) log.debug1('max_memory = %d MB (%d in use) chunks %s', max_memory, mem_now, chunks) log.debug2('j3c_jobs %s', j3c_jobs) if j_only: int3c = wrap_int3c(cell, fused_cell, 'int3c2e', 's2', 1, kptij_lst) else: int3c = wrap_int3c(cell, fused_cell, 'int3c2e', 's1', 1, kptij_lst) idxb = numpy.tril_indices(nao) idxb = (idxb[0] * nao + idxb[1]).astype('i') aux_loc = fused_cell.ao_loc_nr(fused_cell.cart) def gen_int3c(job_id, ish0, ish1): dataname = 'j3c-chunks/%d' % job_id i0 = ao_loc[ish0] i1 = ao_loc[ish1] dii = i1*(i1+1)//2 - i0*(i0+1)//2 dij = (i1 - i0) * nao if j_only: buflen = max(8, int(max_memory*1e6/16/(nkptij*dii+dii))) else: buflen = max(8, int(max_memory*1e6/16/(nkptij*dij+dij))) auxranges = balance_segs(aux_loc[1:]-aux_loc[:-1], buflen) buflen = max([x[2] for x in auxranges]) buf = numpy.empty(nkptij*dij*buflen, dtype=dtype) buf1 = numpy.empty(dij*buflen, dtype=dtype) naux = aux_loc[-1] for kpt_id, kptij in enumerate(kptij_lst): key = '%s/%d' % (dataname, kpt_id) if aosym_s2[kpt_id]: shape = (naux, dii) else: shape = (naux, dij) if gamma_point(kptij): fswap.create_dataset(key, shape, 'f8') else: fswap.create_dataset(key, shape, 'c16') naux0 = 0 for istep, auxrange in enumerate(auxranges): log.alldebug2("aux_e1 job_id %d step %d", job_id, istep) sh0, sh1, nrow = auxrange sub_slice = (ish0, ish1, 0, cell.nbas, sh0, sh1) if j_only: mat = numpy.ndarray((nkptij,dii,nrow), dtype=dtype, buffer=buf) else: mat = numpy.ndarray((nkptij,dij,nrow), dtype=dtype, buffer=buf) mat = int3c(sub_slice, mat) for k, kptij in enumerate(kptij_lst): h5dat = fswap['%s/%d'%(dataname,k)] v = lib.transpose(mat[k], out=buf1) if not j_only and aosym_s2[k]: idy = idxb[i0*(i0+1)//2:i1*(i1+1)//2] - i0 * nao out = numpy.ndarray((nrow,dii), dtype=v.dtype, buffer=mat[k]) v = numpy.take(v, idy, axis=1, out=out) if gamma_point(kptij): h5dat[naux0:naux0+nrow] = v.real else: h5dat[naux0:naux0+nrow] = v naux0 += nrow 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) Gaux = ft_ao.ft_ao(fused_cell, Gv, None, b, gxyz, Gvbase, kpt).T Gaux = fuse(Gaux) Gaux *= mydf.weighted_coulG(kpt, False, mesh) kLR = lib.transpose(numpy.asarray(Gaux.real, order='C')) kLI = lib.transpose(numpy.asarray(Gaux.imag, order='C')) 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 = fuse(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(' blksize (%d,%d)', Gblksize, ncol) 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): 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, idx 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) 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) mpi_df._assemble(mydf, kptij_lst, j3c_jobs, gen_int3c, ft_fuse, cderi_file, fswap, log)
if __name__ == '__main__': from pyscf import gto mol = gto.Mole() mol.atom = '''C 1.3 .2 .3 C .1 .1 1.1 ''' mol.basis = 'ccpvdz' #mol.basis = {'C': [[0, (2.4, .1, .6), (1.0,.8, .4)], [1, (1.1, 1)]]} #mol.basis = {'C': [[0, (2.4, 1)]]} mol.unit = 'B' mol.build(0,0) L = 5. n = 20 a = numpy.diag([L,L,L]) b = scipy.linalg.inv(a) gs = [n,n,n] gxrange = range(gs[0]+1)+range(-gs[0],0) gyrange = range(gs[1]+1)+range(-gs[1],0) gzrange = range(gs[2]+1)+range(-gs[2],0) gxyz = lib.cartesian_prod((gxrange, gyrange, gzrange)) Gv = 2*numpy.pi * numpy.dot(gxyz, b) import time print(time.clock()) print(numpy.linalg.norm(ft_aopair(mol, Gv, None, 's1', b, gxyz, gs)) - 63.0239113778) print(time.clock()) print(numpy.linalg.norm(ft_ao(mol, Gv, None, b, gxyz, gs))-56.8273147065) print(time.clock())
def test_cartesian_prod(self): arrs = (range(3,9), range(4)) cp = lib.cartesian_prod(arrs) for i,x in enumerate(itertools.product(*arrs)): self.assertTrue(numpy.allclose(x,cp[i]))
def _make_j3c(mydf, cell, auxcell, kptij_lst, cderi_file): t1 = (time.clock(), time.time()) log = logger.Logger(mydf.stdout, mydf.verbose) max_memory = max(2000, mydf.max_memory-lib.current_memory()[0]) fused_cell, fuse = fuse_auxcell(mydf, auxcell) # The ideal way to hold the temporary integrals is to store them in the # cderi_file and overwrite them inplace in the second pass. The current # HDF5 library does not have an efficient way to manage free space in # overwriting. It often leads to the cderi_file ~2 times larger than the # necessary size. For now, dumping the DF integral intermediates to a # separated temporary file can avoid this issue. The DF intermediates may # be terribly huge. The temporary file should be placed in the same disk # as cderi_file. swapfile = tempfile.NamedTemporaryFile(dir=os.path.dirname(cderi_file)) fswap = lib.H5TmpFile(swapfile.name) # Unlink swapfile to avoid trash swapfile = None outcore._aux_e2(cell, fused_cell, fswap, 'int3c2e', aosym='s2', kptij_lst=kptij_lst, dataname='j3c-junk', max_memory=max_memory) t1 = log.timer_debug1('3c2e', *t1) nao = cell.nao_nr() naux = auxcell.nao_nr() mesh = mydf.mesh Gv, Gvbase, kws = cell.get_Gv_weights(mesh) b = cell.reciprocal_vectors() gxyz = lib.cartesian_prod([numpy.arange(len(x)) for x in Gvbase]) ngrids = gxyz.shape[0] kptis = kptij_lst[:,0] kptjs = kptij_lst[:,1] kpt_ji = kptjs - kptis uniq_kpts, uniq_index, uniq_inverse = unique(kpt_ji) log.debug('Num uniq kpts %d', len(uniq_kpts)) log.debug2('uniq_kpts %s', uniq_kpts) # j2c ~ (-kpt_ji | kpt_ji) j2c = fused_cell.pbc_intor('int2c2e', hermi=1, kpts=uniq_kpts) max_memory = max(2000, mydf.max_memory - lib.current_memory()[0]) blksize = max(2048, int(max_memory*.5e6/16/fused_cell.nao_nr())) log.debug2('max_memory %s (MB) blocksize %s', max_memory, blksize) for k, kpt in enumerate(uniq_kpts): coulG = mydf.weighted_coulG(kpt, False, mesh) for p0, p1 in lib.prange(0, ngrids, blksize): aoaux = ft_ao.ft_ao(fused_cell, Gv[p0:p1], None, b, gxyz[p0:p1], Gvbase, kpt).T LkR = numpy.asarray(aoaux.real, order='C') LkI = numpy.asarray(aoaux.imag, order='C') aoaux = None if is_zero(kpt): # kpti == kptj j2c[k][naux:] -= lib.ddot(LkR[naux:]*coulG[p0:p1], LkR.T) j2c[k][naux:] -= lib.ddot(LkI[naux:]*coulG[p0:p1], LkI.T) j2c[k][:naux,naux:] = j2c[k][naux:,:naux].T else: j2cR, j2cI = zdotCN(LkR[naux:]*coulG[p0:p1], LkI[naux:]*coulG[p0:p1], LkR.T, LkI.T) j2c[k][naux:] -= j2cR + j2cI * 1j j2c[k][:naux,naux:] = j2c[k][naux:,:naux].T.conj() LkR = LkI = None fswap['j2c/%d'%k] = fuse(fuse(j2c[k]).T).T j2c = coulG = None def cholesky_decomposed_metric(uniq_kptji_id): j2c = numpy.asarray(fswap['j2c/%d'%uniq_kptji_id]) j2c_negative = None try: j2c = scipy.linalg.cholesky(j2c, lower=True) j2ctag = 'CD' except scipy.linalg.LinAlgError as e: #msg =('===================================\n' # 'J-metric not positive definite.\n' # 'It is likely that mesh is not enough.\n' # '===================================') #log.error(msg) #raise scipy.linalg.LinAlgError('\n'.join([str(e), msg])) w, v = scipy.linalg.eigh(j2c) log.debug('DF metric linear dependency for kpt %s', uniq_kptji_id) log.debug('cond = %.4g, drop %d bfns', w[-1]/w[0], numpy.count_nonzero(w<mydf.linear_dep_threshold)) v1 = v[:,w>mydf.linear_dep_threshold].conj().T v1 /= numpy.sqrt(w[w>mydf.linear_dep_threshold]).reshape(-1,1) j2c = v1 if cell.dimension == 2 and cell.low_dim_ft_type != 'inf_vacuum': idx = numpy.where(w < -mydf.linear_dep_threshold)[0] if len(idx) > 0: j2c_negative = (v[:,idx]/numpy.sqrt(-w[idx])).conj().T w = v = None j2ctag = 'eig' return j2c, j2c_negative, j2ctag feri = h5py.File(cderi_file, 'w') feri['j3c-kptij'] = kptij_lst nsegs = len(fswap['j3c-junk/0']) def make_kpt(uniq_kptji_id, cholesky_j2c): kpt = uniq_kpts[uniq_kptji_id] # kpt = kptj - kpti log.debug1('kpt = %s', kpt) adapted_ji_idx = numpy.where(uniq_inverse == uniq_kptji_id)[0] adapted_kptjs = kptjs[adapted_ji_idx] nkptj = len(adapted_kptjs) log.debug1('adapted_ji_idx = %s', adapted_ji_idx) j2c, j2c_negative, j2ctag = cholesky_j2c shls_slice = (auxcell.nbas, fused_cell.nbas) Gaux = ft_ao.ft_ao(fused_cell, Gv, shls_slice, b, gxyz, Gvbase, kpt) wcoulG = mydf.weighted_coulG(kpt, False, mesh) Gaux *= wcoulG.reshape(-1,1) kLR = Gaux.real.copy('C') kLI = Gaux.imag.copy('C') Gaux = None if is_zero(kpt): # kpti == kptj aosym = 's2' nao_pair = nao*(nao+1)//2 if 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] else: aosym = 's1' nao_pair = nao**2 mem_now = lib.current_memory()[0] log.debug2('memory = %s', mem_now) max_memory = max(2000, mydf.max_memory-mem_now) # nkptj for 3c-coulomb arrays plus 1 Lpq array buflen = min(max(int(max_memory*.38e6/16/naux/(nkptj+1)), 1), nao_pair) shranges = _guess_shell_ranges(cell, buflen, aosym) buflen = max([x[2] for x in shranges]) # +1 for a pqkbuf if aosym == 's2': Gblksize = max(16, int(max_memory*.1e6/16/buflen/(nkptj+1))) else: Gblksize = max(16, int(max_memory*.2e6/16/buflen/(nkptj+1))) Gblksize = min(Gblksize, ngrids, 16384) pqkRbuf = numpy.empty(buflen*Gblksize) pqkIbuf = numpy.empty(buflen*Gblksize) # buf for ft_aopair buf = numpy.empty(nkptj*buflen*Gblksize, dtype=numpy.complex128) 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][naux:], 1) lib.dot(kLI[p0:p1].T, pqkI.T, -1, j3cR[k][naux:], 1) if not (is_zero(kpt) and gamma_point(adapted_kptjs[k])): lib.dot(kLR[p0:p1].T, pqkI.T, -1, j3cI[k][naux:], 1) lib.dot(kLI[p0:p1].T, pqkR.T, 1, j3cI[k][naux:], 1) for k, ji in enumerate(adapted_ji_idx): if is_zero(kpt) and gamma_point(adapted_kptjs[k]): v = fuse(j3cR[k]) else: v = fuse(j3cR[k] + j3cI[k] * 1j) if j2ctag == 'CD': v = scipy.linalg.solve_triangular(j2c, v, lower=True, overwrite_b=True) 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) with lib.call_in_background(pw_contract) as compute: col1 = 0 for istep, sh_range in enumerate(shranges): log.debug1('int3c2e [%d/%d], AO [%d:%d], ncol = %d', \ istep+1, len(shranges), *sh_range) bstart, bend, ncol = sh_range col0, col1 = col1, col1+ncol j3cR = [] j3cI = [] for k, idx in enumerate(adapted_ji_idx): v = numpy.vstack([fswap['j3c-junk/%d/%d'%(idx,i)][0,col0:col1].T for i in range(nsegs)]) # vbar is the interaction between the background charge # and the auxiliary basis. 0D, 1D, 2D do not have vbar. if is_zero(kpt) and cell.dimension == 3: for i in numpy.where(vbar != 0)[0]: v[i] -= vbar[i] * ovlp[k][col0:col1] j3cR.append(numpy.asarray(v.real, order='C')) if is_zero(kpt) and gamma_point(adapted_kptjs[k]): j3cI.append(None) else: j3cI.append(numpy.asarray(v.imag, order='C')) v = None compute(istep, sh_range, j3cR, j3cI) for ji in adapted_ji_idx: del(fswap['j3c-junk/%d'%ji]) # Wrapped around boundary and symmetry between k and -k can be used # explicitly for the metric integrals. We consider this symmetry # because it is used in the df_ao2mo module when contracting two 3-index # integral tensors to the 4-index 2e integral tensor. If the symmetry # related k-points are treated separately, the resultant 3-index tensors # may have inconsistent dimension due to the numerial noise when handling # linear dependency of j2c. def conj_j2c(cholesky_j2c): j2c, j2c_negative, j2ctag = cholesky_j2c if j2c_negative is None: return j2c.conj(), None, j2ctag else: return j2c.conj(), j2c_negative.conj(), j2ctag a = cell.lattice_vectors() / (2*numpy.pi) def kconserve_indices(kpt): '''search which (kpts+kpt) satisfies momentum conservation''' kdif = numpy.einsum('wx,ix->wi', a, uniq_kpts + kpt) kdif_int = numpy.rint(kdif) mask = numpy.einsum('wi->i', abs(kdif - kdif_int)) < KPT_DIFF_TOL uniq_kptji_ids = numpy.where(mask)[0] return uniq_kptji_ids done = numpy.zeros(len(uniq_kpts), dtype=bool) for k, kpt in enumerate(uniq_kpts): if done[k]: continue log.debug1('Cholesky decomposition for j2c at kpt %s', k) cholesky_j2c = cholesky_decomposed_metric(k) # The k-point k' which has (k - k') * a = 2n pi. Metric integrals have the # symmetry S = S uniq_kptji_ids = kconserve_indices(-kpt) log.debug1("Symmetry pattern (k - %s)*a= 2n pi", kpt) log.debug1(" make_kpt for uniq_kptji_ids %s", uniq_kptji_ids) for uniq_kptji_id in uniq_kptji_ids: if not done[uniq_kptji_id]: make_kpt(uniq_kptji_id, cholesky_j2c) done[uniq_kptji_ids] = True # The k-point k' which has (k + k') * a = 2n pi. Metric integrals have the # symmetry S = S* uniq_kptji_ids = kconserve_indices(kpt) log.debug1("Symmetry pattern (k + %s)*a= 2n pi", kpt) log.debug1(" make_kpt for %s", uniq_kptji_ids) cholesky_j2c = conj_j2c(cholesky_j2c) for uniq_kptji_id in uniq_kptji_ids: if not done[uniq_kptji_id]: make_kpt(uniq_kptji_id, cholesky_j2c) done[uniq_kptji_ids] = True feri.close()
def _make_j3c(mydf, mol, auxmol): log = logger.Logger(mydf.stdout, mydf.verbose) atm, bas, env, ao_loc = incore._env_and_aoloc('cint3c2e_sph', mol, auxmol) nao = ao_loc[mol.nbas] naux = ao_loc[-1] - nao nao_pair = nao * (nao+1) // 2 cintopt = gto.moleintor.make_cintopt(atm, bas, env, 'cint3c2e_sph') if mydf.approx_sr_level == 0: get_Lpq = _make_Lpq(mydf, mol, auxmol) else: get_Lpq = _make_Lpq_atomic_approx(mydf, mol, auxmol) feri = h5py.File(mydf._cderi) chunks = (min(256,naux), min(256,nao_pair)) # 512K feri.create_dataset('j3c', (naux,nao_pair), 'f8', chunks=chunks) feri.create_dataset('Lpq', (naux,nao_pair), 'f8', chunks=chunks) def save(label, dat, col0, col1): feri[label][:,col0:col1] = dat Gv, Gvbase, kws = non_uniform_kgrids(mydf.gs) nxyz = [i*2 for i in mydf.gs] gxyz = lib.cartesian_prod([range(i) for i in nxyz]) kk = numpy.einsum('ki,ki->k', Gv, Gv) idx = numpy.argsort(kk)[::-1] kk = kk[idx] Gv = Gv[idx] kws = kws[idx] gxyz = gxyz[idx] coulG = .5/numpy.pi**2 * kws / kk aoaux = ft_ao.ft_ao(auxmol, Gv, None, Gvbase, gxyz, nxyz) kLR = numpy.asarray(aoaux.real, order='C') kLI = numpy.asarray(aoaux.imag, order='C') j2c = auxmol.intor('cint2c2e_sph', hermi=1).T # .T to C-ordr lib.dot(kLR.T*coulG, kLR, -1, j2c, 1) lib.dot(kLI.T*coulG, kLI, -1, j2c, 1) kLR *= coulG.reshape(-1,1) kLI *= coulG.reshape(-1,1) aoaux = coulG = kk = kws = idx = None max_memory = max(2000, mydf.max_memory-lib.current_memory()[0]) buflen = min(max(int(max_memory*.3*1e6/8/naux), 1), nao_pair) shranges = outcore._guess_shell_ranges(mol, buflen, 's2ij') buflen = max([x[2] for x in shranges]) blksize = max(16, int(max_memory*.15*1e6/16/buflen)) pqkbuf = numpy.empty(buflen*blksize) bufs1 = numpy.empty((buflen*naux)) # bufs2 holds either Lpq and ft_aopair bufs2 = numpy.empty(max(buflen*(naux+1),buflen*blksize*2)) # *2 for cmplx col1 = 0 for istep, sh_range in enumerate(shranges): log.debug('int3c2e [%d/%d], AO [%d:%d], ncol = %d', \ istep+1, len(shranges), *sh_range) bstart, bend, ncol = sh_range col0, col1 = col1, col1+ncol shls_slice = (bstart, bend, 0, bend, mol.nbas, mol.nbas+auxmol.nbas) Lpq = get_Lpq(shls_slice, col0, col1, bufs2) save('Lpq', Lpq, col0, col1) j3c = _ri.nr_auxe2('cint3c2e_sph', atm, bas, env, shls_slice, ao_loc, 's2ij', 1, cintopt, bufs1) j3c = j3c.T # -> (L|pq) in C-order lib.dot(j2c, Lpq, -.5, j3c, 1) Lpq = None for p0, p1 in lib.prange(0, Gv.shape[0], blksize): aoao = ft_ao.ft_aopair(mol, Gv[p0:p1], shls_slice[:4], 's2', Gvbase, gxyz[p0:p1], nxyz, buf=bufs2) nG = p1 - p0 pqkR = numpy.ndarray((ncol,nG), buffer=pqkbuf) pqkR[:] = aoao.real.T lib.dot(kLR[p0:p1].T, pqkR.T, -1, j3c, 1) pqkI = numpy.ndarray((ncol,nG), buffer=pqkbuf) pqkI[:] = aoao.imag.T lib.dot(kLI[p0:p1].T, pqkI.T, -1, j3c, 1) aoao = aoaux = None save('j3c', j3c, col0, col1) feri.close()
def ewald(cell, ew_eta=None, ew_cut=None): '''Perform real (R) and reciprocal (G) space Ewald sum for the energy. Formulation of Martin, App. F2. Returns: float The Ewald energy consisting of overlap, self, and G-space sum. See Also: pyscf.pbc.gto.get_ewald_params ''' if ew_eta is None: ew_eta = cell.ew_eta if ew_cut is None: ew_cut = cell.ew_cut chargs = cell.atom_charges() coords = cell.atom_coords() ewovrl = 0. # set up real-space lattice indices [-ewcut ... ewcut] ewxrange = np.arange(-ew_cut[0],ew_cut[0]+1) ewyrange = np.arange(-ew_cut[1],ew_cut[1]+1) ewzrange = np.arange(-ew_cut[2],ew_cut[2]+1) ewxyz = lib.cartesian_prod((ewxrange,ewyrange,ewzrange)).T nx = len(ewxrange) ny = len(ewyrange) nz = len(ewzrange) Lall = np.einsum('ij,jk->ik', cell._h, ewxyz).reshape(3,nx,ny,nz) #exclude the point where Lall == 0 Lall[:,ew_cut[0],ew_cut[1],ew_cut[2]] = 1e200 Lall = Lall.reshape(3,nx*ny*nz) Lall = Lall.T for ia in range(cell.natm): qi = chargs[ia] ri = coords[ia] for ja in range(ia): qj = chargs[ja] rj = coords[ja] r = np.linalg.norm(ri-rj) ewovrl += 2 * qi * qj / r * scipy.special.erfc(ew_eta * r) for ia in range(cell.natm): qi = chargs[ia] ri = coords[ia] for ja in range(cell.natm): qj = chargs[ja] rj = coords[ja] r1 = ri-rj + Lall r = np.sqrt(np.einsum('ji,ji->j', r1, r1)) ewovrl += (qi * qj / r * scipy.special.erfc(ew_eta * r)).sum() ewovrl *= 0.5 # last line of Eq. (F.5) in Martin ewself = -1./2. * np.dot(chargs,chargs) * 2 * ew_eta / np.sqrt(np.pi) ewself += -1./2. * np.sum(chargs)**2 * np.pi/(ew_eta**2 * cell.vol) # g-space sum (using g grid) (Eq. (F.6) in Martin, but note errors as below) SI = cell.get_SI() ZSI = np.einsum("i,ij->j", chargs, SI) # Eq. (F.6) in Martin is off by a factor of 2, the # exponent is wrong (8->4) and the square is in the wrong place # # Formula should be # 1/2 * 4\pi / Omega \sum_I \sum_{G\neq 0} |ZS_I(G)|^2 \exp[-|G|^2/4\eta^2] # where # ZS_I(G) = \sum_a Z_a exp (i G.R_a) # See also Eq. (32) of ewald.pdf at # http://www.fisica.uniud.it/~giannozz/public/ewald.pdf coulG = pbctools.get_coulG(cell) absG2 = np.einsum('gi,gi->g', cell.Gv, cell.Gv) ZSIG2 = np.abs(ZSI)**2 expG2 = np.exp(-absG2/(4*ew_eta**2)) JexpG2 = coulG*expG2 ewgI = np.dot(ZSIG2,JexpG2) ewg = .5*np.sum(ewgI) ewg /= cell.vol #log.debug('Ewald components = %.15g, %.15g, %.15g', ewovrl, ewself, ewg) return ewovrl + ewself + ewg