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 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
