def get_wannier(self, supercell = [1,1,1], grid = [50,50,50]):
'''
Evaluate the MLWF using a periodic grid
'''
grids_coor, weights = periodic_grid(self.cell, grid, supercell = [1,1,1], order = 'C')
kpts = self.cell.get_abs_kpts(self.kpt_latt_loc)
ao_kpts = np.asarray([numint.eval_ao(self.cell, grids_coor, kpt = kpt) for kpt in kpts])
u_mo = []
for k_id in range(self.num_kpts_loc):
mo_included = self.mo_coeff_kpts[k_id][:,self.band_included_list]
mo_in_window = self.lwindow[k_id]
C_opt = mo_included[:,mo_in_window].dot(self.U_matrix_opt[k_id].T)
C_tildle = C_opt.dot(self.U_matrix[k_id].T)
kpt = kpts[k_id]
ao = numint.eval_ao(self.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)
WF0 = libwannier90.get_WF0s(self.kpt_latt_loc.shape[0],self.kpt_latt_loc, supercell, grid, u_mo)
# Fix the global phase following the pw2wannier90 procedure
max_index = (WF0*WF0.conj()).real.argmax(axis=0)
norm_wfs = np.diag(WF0[max_index,:])
norm_wfs = norm_wfs/np.absolute(norm_wfs)
WF0 = WF0/norm_wfs/self.num_kpts_loc
# Check the 'reality' following the pw2wannier90 procedure
for WF_id in range(self.num_wann_loc):
ratio_max = np.abs(WF0[np.abs(WF0[:,WF_id].real) >= 0.01,WF_id].imag/WF0[np.abs(WF0[:,WF_id].real) >= 0.01,WF_id].real).max(axis=0)
print('The maximum imag/real for wannier function ', WF_id,' : ', ratio_max)
return WF0
def aux_e2_grid(cell, auxcell, grids=None):
'''3-center AO integrals (ij|L), where L is the auxiliary basis.
Implements double summation over lattice vectors: \sum_{lm} (i[l]j[m]|L[0]).
'''
if grids is None:
grids = gen_grid.BeckeGrids(cell)
grids.build_()
elif grids.weights is None:
raise RuntimeError('grids is not initialized')
ao = numpy.asarray(numint.eval_ao(cell, grids.coords).real, order='C')
auxao = numpy.asarray(numint.eval_ao(auxcell, grids.coords).real, order='C')
if isinstance(grids.weights, numpy.ndarray):
auxao *= grids.weights.reshape(-1,1)
else:
auxao *= grids.weights
ng, nao = ao.shape
naoaux = auxao.shape[1]
#aux_e2 = numpy.einsum('ri,rj,rk',ao,ao,auxao)
aux_e2 = numpy.zeros((nao*nao,naoaux))
for p0, p1 in prange(0, ng, 240):
tmp = numpy.einsum('ri,rj->rij', ao[p0:p1], ao[p0:p1])
pyscf.lib.dot(tmp.reshape(p1-p0,-1).T, auxao[p0:p1], 1, aux_e2, 1)
return aux_e2.reshape(nao,nao,-1)
def get_jk(mf, cell, dm, hermi=1, vhfopt=None, kpt=np.zeros(3), kpts_band=None):
dm = np.asarray(dm)
nao = dm.shape[-1]
coords = gen_grid.gen_uniform_grids(cell)
if kpts_band is None:
kpt1 = kpt2 = kpt
aoR_k1 = aoR_k2 = numint.eval_ao(cell, coords, kpt)
else:
kpt1 = kpts_band
kpt2 = kpt
aoR_k1 = numint.eval_ao(cell, coords, kpt1)
aoR_k2 = numint.eval_ao(cell, coords, kpt2)
vkR_k1k2 = get_vkR(mf, cell, aoR_k1, aoR_k2, kpt1, kpt2)
ngrids, nao = aoR_k1.shape
def contract(dm):
vjR_k2 = get_vjR(cell, dm, aoR_k2)
vj = (cell.vol/ngrids) * np.dot(aoR_k1.T.conj(), vjR_k2.reshape(-1,1)*aoR_k1)
#:vk = (cell.vol/ngrids) * np.einsum('rs,Rp,Rqs,Rr->pq', dm, aoR_k1.conj(),
#: vkR_k1k2, aoR_k2)
aoR_dm_k2 = np.dot(aoR_k2, dm)
tmp_Rq = np.einsum('Rqs,Rs->Rq', vkR_k1k2, aoR_dm_k2)
vk = (cell.vol/ngrids) * np.dot(aoR_k1.T.conj(), tmp_Rq)
return vj, vk
if dm.ndim == 2:
vj, vk = contract(dm)
else:
jk = [contract(x) for x in dm.reshape(-1,nao,nao)]
vj = lib.asarray([x[0] for x in jk])
vk = lib.asarray([x[1] for x in jk])
return vj.reshape(dm.shape), vk.reshape(dm.shape)
def get_jk(mf, cell, dm, hermi=1, vhfopt=None, kpt=np.zeros(3), kpt_band=None):
dm = np.asarray(dm)
nao = dm.shape[-1]
coords = gen_grid.gen_uniform_grids(cell)
if kpt_band is None:
kpt1 = kpt2 = kpt
aoR_k1 = aoR_k2 = numint.eval_ao(cell, coords, kpt)
else:
kpt1 = kpt_band
kpt2 = kpt
aoR_k1 = numint.eval_ao(cell, coords, kpt1)
aoR_k2 = numint.eval_ao(cell, coords, kpt2)
vkR_k1k2 = get_vkR(mf, cell, aoR_k1, aoR_k2, kpt1, kpt2)
ngs, nao = aoR_k1.shape
def contract(dm):
vjR_k2 = get_vjR(cell, dm, aoR_k2)
vj = (cell.vol/ngs) * np.dot(aoR_k1.T.conj(), vjR_k2.reshape(-1,1)*aoR_k1)
#:vk = (cell.vol/ngs) * np.einsum('rs,Rp,Rqs,Rr->pq', dm, aoR_k1.conj(),
#: vkR_k1k2, aoR_k2)
aoR_dm_k2 = np.dot(aoR_k2, dm)
tmp_Rq = np.einsum('Rqs,Rs->Rq', vkR_k1k2, aoR_dm_k2)
vk = (cell.vol/ngs) * np.dot(aoR_k1.T.conj(), tmp_Rq)
return vj, vk
if dm.ndim == 2:
vj, vk = contract(dm)
else:
jk = [contract(x) for x in dm.reshape(-1,nao,nao)]
vj = lib.asarray([x[0] for x in jk])
vk = lib.asarray([x[1] for x in jk])
return vj.reshape(dm.shape), vk.reshape(dm.shape)
def get_j(cell, dm, hermi=1, vhfopt=None, kpt=np.zeros(3), kpts_band=None):
dm = np.asarray(dm)
nao = dm.shape[-1]
coords = gen_grid.gen_uniform_grids(cell)
if kpts_band is None:
kpt1 = kpt2 = kpt
aoR_k1 = aoR_k2 = numint.eval_ao(cell, coords, kpt)
else:
kpt1 = kpts_band
kpt2 = kpt
aoR_k1 = numint.eval_ao(cell, coords, kpt1)
aoR_k2 = numint.eval_ao(cell, coords, kpt2)
ngs, nao = aoR_k1.shape
def contract(dm):
vjR_k2 = get_vjR(cell, dm, aoR_k2)
vj = (cell.vol / ngs) * np.dot(aoR_k1.T.conj(),
vjR_k2.reshape(-1, 1) * aoR_k1)
return vj
if dm.ndim == 2:
vj = contract(dm)
else:
vj = lib.asarray([contract(x) for x in dm.reshape(-1, nao, nao)])
return vj.reshape(dm.shape)
def get_j(cell, dm, hermi=1, vhfopt=None, kpt=np.zeros(3), kpt_band=None):
dm = np.asarray(dm)
nao = dm.shape[-1]
coords = gen_grid.gen_uniform_grids(cell)
if kpt_band is None:
kpt1 = kpt2 = kpt
aoR_k1 = aoR_k2 = numint.eval_ao(cell, coords, kpt)
else:
kpt1 = kpt_band
kpt2 = kpt
aoR_k1 = numint.eval_ao(cell, coords, kpt1)
aoR_k2 = numint.eval_ao(cell, coords, kpt2)
ngs, nao = aoR_k1.shape
def contract(dm):
vjR_k2 = get_vjR(cell, dm, aoR_k2)
vj = (cell.vol/ngs) * np.dot(aoR_k1.T.conj(), vjR_k2.reshape(-1,1)*aoR_k1)
return vj
if dm.ndim == 2:
vj = contract(dm)
else:
vj = lib.asarray([contract(x) for x in dm.reshape(-1,nao,nao)])
return vj.reshape(dm.shape)
def get_wannier(self, supercell = [1,1,1], grid = [50,50,50]):
'''
Evaluate the MLWF using a periodic grid
'''
grids_coor, weights = periodic_grid(self.cell, grid, supercell = [1,1,1], order = 'C')
kpts = self.cell.get_abs_kpts(self.kpt_latt_loc)
ao_kpts = np.asarray([numint.eval_ao(self.cell, grids_coor, kpt = kpt) for kpt in kpts])
u_mo = []
for k_id in range(self.num_kpts_loc):
mo_included = self.mo_coeff_kpts[k_id][:,self.band_included_list]
mo_in_window = self.lwindow[k_id]
C_opt = mo_included[:,mo_in_window].dot(self.U_matrix_opt[k_id].T)
C_tildle = C_opt.dot(self.U_matrix[k_id].T)
kpt = kpts[k_id]
ao = numint.eval_ao(self.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)
WF0 = libwannier90.get_WF0s(self.kpt_latt_loc.shape[0],self.kpt_latt_loc, supercell, grid, u_mo)
# Fix the global phase following the pw2wannier90 procedure
max_index = (WF0*WF0.conj()).real.argmax(axis=0)
norm_wfs = np.diag(WF0[max_index,:])
norm_wfs = norm_wfs/np.absolute(norm_wfs)
WF0 = WF0/norm_wfs/self.num_kpts_loc
# Check the 'reality' following the pw2wannier90 procedure
for WF_id in range(self.num_wann_loc):
ratio_max = np.abs(WF0[np.abs(WF0[:,WF_id].real) >= 0.01,WF_id].imag/WF0[np.abs(WF0[:,WF_id].real) >= 0.01,WF_id].real).max(axis=0)
print('The maximum imag/real for wannier function ', WF_id,' : ', ratio_max)
return WF0
def get_mo_pairs_G_old(cell, mo_coeffs, kpts=None, q=None):
'''Calculate forward (G|ij) and "inverse" (ij|G) FFT of all MO pairs.
TODO: - Implement simplifications for real orbitals.
Args:
mo_coeff: length-2 list of (nao,nmo) ndarrays
The two sets of MO coefficients to use in calculating the
product |ij).
Returns:
mo_pairs_G, mo_pairs_invG : (ngrids, nmoi*nmoj) ndarray
The FFTs of the real-space MO pairs.
'''
coords = gen_uniform_grids(cell)
if kpts is None:
q = np.zeros(3)
aoR = eval_ao(cell, coords)
ngrids = aoR.shape[0]
if np.array_equal(mo_coeffs[0], mo_coeffs[1]):
nmoi = nmoj = mo_coeffs[0].shape[1]
moiR = mojR = einsum('ri,ia->ra', aoR, mo_coeffs[0])
else:
nmoi = mo_coeffs[0].shape[1]
nmoj = mo_coeffs[1].shape[1]
moiR = einsum('ri,ia->ra', aoR, mo_coeffs[0])
mojR = einsum('ri,ia->ra', aoR, mo_coeffs[1])
else:
if q is None:
q = kpts[1] - kpts[0]
aoR_ki = eval_ao(cell, coords, kpt=kpts[0])
aoR_kj = eval_ao(cell, coords, kpt=kpts[1])
ngrids = aoR_ki.shape[0]
nmoi = mo_coeffs[0].shape[1]
nmoj = mo_coeffs[1].shape[1]
moiR = einsum('ri,ia->ra', aoR_ki, mo_coeffs[0])
mojR = einsum('ri,ia->ra', aoR_kj, mo_coeffs[1])
mo_pairs_R = np.einsum('ri,rj->rij', np.conj(moiR), mojR)
mo_pairs_G = np.zeros([ngrids, nmoi * nmoj], np.complex128)
mo_pairs_invG = np.zeros([ngrids, nmoi * nmoj], np.complex128)
fac = np.exp(-1j * np.dot(coords, q))
for i in range(nmoi):
for j in range(nmoj):
mo_pairs_G[:, i * nmoj + j] = tools.fftk(mo_pairs_R[:, i, j],
cell.mesh, fac)
mo_pairs_invG[:, i * nmoj + j] = np.conj(
tools.fftk(np.conj(mo_pairs_R[:, i, j]), cell.mesh,
fac.conj()))
return mo_pairs_G, mo_pairs_invG
def get_mo_pairs_G_old(cell, mo_coeffs, kpts=None, q=None):
'''Calculate forward (G|ij) and "inverse" (ij|G) FFT of all MO pairs.
TODO: - Implement simplifications for real orbitals.
Args:
mo_coeff: length-2 list of (nao,nmo) ndarrays
The two sets of MO coefficients to use in calculating the
product |ij).
Returns:
mo_pairs_G, mo_pairs_invG : (ngs, nmoi*nmoj) ndarray
The FFTs of the real-space MO pairs.
'''
coords = gen_uniform_grids(cell)
if kpts is None:
q = np.zeros(3)
aoR = eval_ao(cell, coords)
ngs = aoR.shape[0]
if np.array_equal(mo_coeffs[0], mo_coeffs[1]):
nmoi = nmoj = mo_coeffs[0].shape[1]
moiR = mojR = einsum('ri,ia->ra', aoR, mo_coeffs[0])
else:
nmoi = mo_coeffs[0].shape[1]
nmoj = mo_coeffs[1].shape[1]
moiR = einsum('ri,ia->ra', aoR, mo_coeffs[0])
mojR = einsum('ri,ia->ra', aoR, mo_coeffs[1])
else:
if q is None:
q = kpts[1]-kpts[0]
aoR_ki = eval_ao(cell, coords, kpt=kpts[0])
aoR_kj = eval_ao(cell, coords, kpt=kpts[1])
ngs = aoR_ki.shape[0]
nmoi = mo_coeffs[0].shape[1]
nmoj = mo_coeffs[1].shape[1]
moiR = einsum('ri,ia->ra', aoR_ki, mo_coeffs[0])
mojR = einsum('ri,ia->ra', aoR_kj, mo_coeffs[1])
mo_pairs_R = np.einsum('ri,rj->rij', np.conj(moiR), mojR)
mo_pairs_G = np.zeros([ngs,nmoi*nmoj], np.complex128)
mo_pairs_invG = np.zeros([ngs,nmoi*nmoj], np.complex128)
fac = np.exp(-1j*np.dot(coords, q))
for i in xrange(nmoi):
for j in xrange(nmoj):
mo_pairs_G[:,i*nmoj+j] = tools.fftk(mo_pairs_R[:,i,j], cell.gs, fac)
mo_pairs_invG[:,i*nmoj+j] = np.conj(tools.fftk(np.conj(mo_pairs_R[:,i,j]), cell.gs,
fac.conj()))
return mo_pairs_G, mo_pairs_invG
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 get_j_kpts(mf, cell, dm_kpts, kpts, kpt_band=None):
coords = gen_grid.gen_uniform_grids(cell)
nkpts = len(kpts)
ngs = len(coords)
dm_kpts = np.asarray(dm_kpts)
nao = dm_kpts.shape[-1]
ni = numint._KNumInt(kpts)
aoR_kpts = ni.eval_ao(cell, coords, kpts)
if kpt_band is not None:
aoR_kband = numint.eval_ao(cell, coords, kpt_band)
dms = dm_kpts.reshape(-1,nkpts,nao,nao)
nset = dms.shape[0]
vjR = [get_vjR(cell, dms[i], aoR_kpts) for i in range(nset)]
if kpt_band is not None:
vj_kpts = [cell.vol/ngs * lib.dot(aoR_kband.T.conj()*vjR[i], aoR_kband)
for i in range(nset)]
if dm_kpts.ndim == 3: # One set of dm_kpts for KRHF
vj_kpts = vj_kpts[0]
return lib.asarray(vj_kpts)
else:
vj_kpts = []
for i in range(nset):
vj = [cell.vol/ngs * lib.dot(aoR_k.T.conj()*vjR[i], aoR_k)
for aoR_k in aoR_kpts]
vj_kpts.append(lib.asarray(vj))
return lib.asarray(vj_kpts).reshape(dm_kpts.shape)
def get_j_kpts(mf, cell, dm_kpts, kpts, kpts_band=None):
coords = gen_grid.gen_uniform_grids(cell)
nkpts = len(kpts)
ngs = len(coords)
dm_kpts = np.asarray(dm_kpts)
nao = dm_kpts.shape[-1]
ni = numint._KNumInt(kpts)
aoR_kpts = ni.eval_ao(cell, coords, kpts)
if kpts_band is not None:
aoR_kband = numint.eval_ao(cell, coords, kpts_band)
dms = dm_kpts.reshape(-1, nkpts, nao, nao)
nset = dms.shape[0]
vjR = [get_vjR(cell, dms[i], aoR_kpts) for i in range(nset)]
if kpts_band is not None:
vj_kpts = [
cell.vol / ngs * lib.dot(aoR_kband.T.conj() * vjR[i], aoR_kband)
for i in range(nset)
]
if dm_kpts.ndim == 3: # One set of dm_kpts for KRHF
vj_kpts = vj_kpts[0]
return lib.asarray(vj_kpts)
else:
vj_kpts = []
for i in range(nset):
vj = [
cell.vol / ngs * lib.dot(aoR_k.T.conj() * vjR[i], aoR_k)
for aoR_k in aoR_kpts
]
vj_kpts.append(lib.asarray(vj))
return lib.asarray(vj_kpts).reshape(dm_kpts.shape)
def get_wannier(self, grid=[50, 50, 50]):
'''
Evaluate the MLWF using a general grid
'''
grids_coor = general_grid(self.cell, grid)
WFs = 0
for k_id in range(self.num_kpts_loc): #self.num_kpts_loc
kpt = self.cell.get_abs_kpts(self.kpt_latt_loc[k_id])
ao = numint.eval_ao(self.cell, grids_coor, kpt=kpt)
mo_included = np.dot(
self.U[k_id],
self.mo_coeff_kpts[k_id])[:, self.band_included_list]
mo_in_window = self.lwindow[k_id]
C_opt = mo_included[:, mo_in_window].dot(self.U_matrix_opt[k_id].T)
C_tildle = C_opt.dot(self.U_matrix[k_id].T)
WFs = WFs + np.einsum('xi,in->xn', ao, C_tildle, optimize=True)
# Fix the global phase following the pw2wannier90 procedure, todo: why?
max_index = (WFs * WFs.conj()).real.argmax(axis=0)
norm_wfs = np.diag(WFs[max_index, :])
norm_wfs = norm_wfs / np.absolute(norm_wfs)
WFs = WFs / norm_wfs / self.num_kpts_loc
# Check the 'reality' following the pw2wannier90 procedure
for WF_id in range(self.num_wann_loc):
ratio_max = np.abs(
WFs[(WFs[:, WF_id].real > 0.01), WF_id].imag /
WFs[(WFs[:, WF_id].real > 0.01), WF_id].real).max(axis=0)
print('The maximum imag/real for wannier function ', WF_id, ' : ',
ratio_max)
return WFs
def export_unk(self, grid=[50, 50, 50]):
'''
Export the periodic part of BF in a real space grid for plotting with wannier90
'''
from scipy.io import FortranFile
grids_coor = general_grid(self.cell, grid)
for k_id in range(self.num_kpts_loc):
kpt = self.cell.get_abs_kpts(self.kpt_latt_loc[k_id])
ao = numint.eval_ao(self.cell, grids_coor, kpt=kpt)
u_ao = np.einsum('x,xi->xi',
np.exp(-1j * np.dot(grids_coor, kpt)),
ao,
optimize=True)
unk_file = FortranFile('UNK0000' + str(k_id + 1) + '.1', 'w')
unk_file.write_record(
np.asarray(
[grid[0], grid[1], grid[2], k_id + 1, self.num_bands_loc],
dtype=np.int32))
mo_included = np.dot(
self.U[k_id],
self.mo_coeff_kpts[k_id])[:, self.band_included_list]
u_mo = np.einsum('xi,in->xn', u_ao, mo_included, optimize=True)
for band in range(len(self.band_included_list)):
unk_file.write_record(
np.asarray(u_mo[:, band], dtype=np.complex))
unk_file.close()
def get_A_mat(self):
r'''
Construct the projection matrix: A_{m,n}^{\mathbf{k}}
Equation (62) in MV, Phys. Rev. B 56, 12847 or equation (22) in SMV, Phys. Rev. B 65, 035109
'''
A_matrix_loc = np.empty([self.num_kpts_loc, self.num_wann_loc, self.num_bands_loc], dtype = np.complex)
if self.use_bloch_phases == True:
for k_id in range(self.num_kpts_loc):
Amn = np.zeros([self.num_wann_loc, self.num_bands_loc])
np.fill_diagonal(Amn, 1)
A_matrix_loc[k_id,:,:] = Amn
else:
grids = gen_grid.UniformGrids(self.cell)
grids.build()
for k_id in range(self.num_kpts_loc):
kpt = self.cell.get_abs_kpts(self.kpt_latt_loc[k_id])
ao = numint.eval_ao(self.cell, grids.coords, kpt = kpt)
for ith_wann in range(self.num_wann_loc):
frac_site = self.proj_site[ith_wann]
abs_site = frac_site.dot(self.real_lattice_loc) / param.BOHR
l = self.proj_l[ith_wann]
mr = self.proj_m[ith_wann]
r = self.proj_radial[ith_wann]
zona = self.proj_zona[ith_wann]
x_axis = self.proj_x[ith_wann]
z_axis = self.proj_z[ith_wann]
gr = g_r(grids.coords, abs_site, l, mr, r, zona, x_axis, z_axis, unit = 'B')
C = np.dot(self.U[k_id], self.mo_coeff_kpts[k_id])[:,self.band_included_list]
s_ao = np.einsum('i,iu,i->u', grids.weights, ao.conj(), gr, optimize = True)
A_matrix_loc[k_id,ith_wann,:] = np.einsum('um,u->m', C, s_ao, optimize = True)
return A_matrix_loc
def get_A_mat(self):
'''
Construct the projection matrix: A_{m,n}^{\mathbf{k}}
Equation (62) in MV, Phys. Rev. B 56, 12847 or equation (22) in SMV, Phys. Rev. B 65, 035109
'''
A_matrix_loc = np.empty(
[self.num_kpts_loc, self.num_wann_loc, self.num_bands_loc],
dtype=np.complex128)
if self.use_bloch_phases == True:
Amn = np.zeros([self.num_wann_loc, self.num_bands_loc])
np.fill_diagonal(Amn, 1)
A_matrix_loc[:, :, :] = Amn
else:
from pyscf.dft import numint, gen_grid
grids = gen_grid.Grids(self.cell).build()
coords = grids.coords
weights = grids.weights
for ith_wann in range(self.num_wann_loc):
frac_site = self.proj_site[ith_wann]
abs_site = frac_site.dot(self.real_lattice_loc) / param.BOHR
l = self.proj_l[ith_wann]
mr = self.proj_m[ith_wann]
r = self.proj_radial[ith_wann]
zona = self.proj_zona[ith_wann]
x_axis = self.proj_x[ith_wann]
z_axis = self.proj_z[ith_wann]
gr = g_r(coords,
abs_site,
l,
mr,
r,
zona,
x_axis,
z_axis,
unit='B')
ao_L0 = numint.eval_ao(self.cell, coords)
s_aoL0_g = np.einsum('i,i,iv->v',
weights,
gr,
ao_L0,
optimize=True)
for k_id in range(self.num_kpts_loc):
kpt = self.cell.get_abs_kpts(self.kpt_latt_loc[k_id])
mo_included = self.mo_coeff_kpts[k_id][:, self.
band_included_list]
s_kpt = self.cell.pbc_intor('int1e_ovlp',
hermi=1,
kpts=kpt,
pbcopt=lib.c_null_ptr())
s_ao = np.einsum('uv,v->u', s_kpt, s_aoL0_g, optimize=True)
A_matrix_loc[k_id, ith_wann, :] = np.einsum(
'v,vu,um->m',
s_aoL0_g,
s_kpt,
mo_included,
optimize=True).conj()
return A_matrix_loc
def export_unk(self, grid=[50, 50, 50]):
'''
Export the periodic part of BF in a real space grid for plotting with wannier90
'''
from scipy.io import FortranFile
grids_coor, weights = periodic_grid(self.cell, grid, order='F')
for k_id in range(self.num_kpts_loc):
spin = '.1'
if self.spin_up != None and self.spin_up == False: spin = '.2'
kpt = self.cell.get_abs_kpts(self.kpt_latt_loc[k_id])
ao = numint.eval_ao(self.cell, grids_coor, kpt=kpt)
u_ao = np.einsum('x,xi->xi',
np.exp(-1j * np.dot(grids_coor, kpt)),
ao,
optimize=True)
unk_file = FortranFile('UNK' + "%05d" % (k_id + 1) + spin, 'w')
unk_file.write_record(
np.asarray(
[grid[0], grid[1], grid[2], k_id + 1, self.num_bands_loc],
dtype=np.int32))
mo_included = self.mo_coeff_kpts[k_id][:, self.band_included_list]
u_mo = np.einsum('xi,in->xn', u_ao, mo_included, optimize=True)
for band in range(len(self.band_included_list)):
unk_file.write_record(
np.asarray(u_mo[:, band], dtype=np.complex128))
unk_file.close()
def get_jk_kpts(mf, cell, dm_kpts, kpts, kpts_band=None):
coords = gen_grid.gen_uniform_grids(cell)
nkpts = len(kpts)
ngrids = len(coords)
dm_kpts = np.asarray(dm_kpts)
nao = dm_kpts.shape[-1]
dms = dm_kpts.reshape(-1,nkpts,nao,nao)
nset = dms.shape[0]
ni = numint.KNumInt(kpts)
aoR_kpts = ni.eval_ao(cell, coords, kpts)
if kpts_band is not None:
aoR_kband = numint.eval_ao(cell, coords, kpts_band)
# J
vjR = [get_vjR_kpts(cell, dms[i], aoR_kpts) for i in range(nset)]
if kpts_band is not None:
vj_kpts = [cell.vol/ngrids * lib.dot(aoR_kband.T.conj()*vjR[i], aoR_kband)
for i in range(nset)]
else:
vj_kpts = []
for i in range(nset):
vj = [cell.vol/ngrids * lib.dot(aoR_k.T.conj()*vjR[i], aoR_k)
for aoR_k in aoR_kpts]
vj_kpts.append(lib.asarray(vj))
vj_kpts = lib.asarray(vj_kpts)
vjR = None
# K
weight = 1./nkpts * (cell.vol/ngrids)
vk_kpts = np.zeros_like(vj_kpts)
if kpts_band is not None:
for k2, kpt2 in enumerate(kpts):
aoR_dms = [lib.dot(aoR_kpts[k2], dms[i,k2]) for i in range(nset)]
vkR_k1k2 = get_vkR(mf, cell, aoR_kband, aoR_kpts[k2],
kpts_band, kpt2)
#:vk_kpts = 1./nkpts * (cell.vol/ngrids) * np.einsum('rs,Rp,Rqs,Rr->pq',
#: dm_kpts[k2], aoR_kband.conj(),
#: vkR_k1k2, aoR_kpts[k2])
for i in range(nset):
tmp_Rq = np.einsum('Rqs,Rs->Rq', vkR_k1k2, aoR_dms[i])
vk_kpts[i] += weight * lib.dot(aoR_kband.T.conj(), tmp_Rq)
vkR_k1k2 = None
if dm_kpts.ndim == 3:
vj_kpts = vj_kpts[0]
vk_kpts = vk_kpts[0]
return lib.asarray(vj_kpts), lib.asarray(vk_kpts)
else:
for k2, kpt2 in enumerate(kpts):
aoR_dms = [lib.dot(aoR_kpts[k2], dms[i,k2]) for i in range(nset)]
for k1, kpt1 in enumerate(kpts):
vkR_k1k2 = get_vkR(mf, cell, aoR_kpts[k1], aoR_kpts[k2],
kpt1, kpt2)
for i in range(nset):
tmp_Rq = np.einsum('Rqs,Rs->Rq', vkR_k1k2, aoR_dms[i])
vk_kpts[i,k1] += weight * lib.dot(aoR_kpts[k1].T.conj(), tmp_Rq)
vkR_k1k2 = None
return vj_kpts.reshape(dm_kpts.shape), vk_kpts.reshape(dm_kpts.shape)
def get_jk_kpts(mf, cell, dm_kpts, kpts, kpt_band=None):
coords = gen_grid.gen_uniform_grids(cell)
nkpts = len(kpts)
ngs = len(coords)
dm_kpts = np.asarray(dm_kpts)
nao = dm_kpts.shape[-1]
dms = dm_kpts.reshape(-1,nkpts,nao,nao)
nset = dms.shape[0]
ni = numint._KNumInt(kpts)
aoR_kpts = ni.eval_ao(cell, coords, kpts)
if kpt_band is not None:
aoR_kband = numint.eval_ao(cell, coords, kpt_band)
# J
vjR = [get_vjR_kpts(cell, dms[i], aoR_kpts) for i in range(nset)]
if kpt_band is not None:
vj_kpts = [cell.vol/ngs * lib.dot(aoR_kband.T.conj()*vjR[i], aoR_kband)
for i in range(nset)]
else:
vj_kpts = []
for i in range(nset):
vj = [cell.vol/ngs * lib.dot(aoR_k.T.conj()*vjR[i], aoR_k)
for aoR_k in aoR_kpts]
vj_kpts.append(lib.asarray(vj))
vj_kpts = lib.asarray(vj_kpts)
vjR = None
# K
weight = 1./nkpts * (cell.vol/ngs)
vk_kpts = np.zeros_like(vj_kpts)
if kpt_band is not None:
for k2, kpt2 in enumerate(kpts):
aoR_dms = [lib.dot(aoR_kpts[k2], dms[i,k2]) for i in range(nset)]
vkR_k1k2 = get_vkR(mf, cell, aoR_kband, aoR_kpts[k2],
kpt_band, kpt2)
#:vk_kpts = 1./nkpts * (cell.vol/ngs) * np.einsum('rs,Rp,Rqs,Rr->pq',
#: dm_kpts[k2], aoR_kband.conj(),
#: vkR_k1k2, aoR_kpts[k2])
for i in range(nset):
tmp_Rq = np.einsum('Rqs,Rs->Rq', vkR_k1k2, aoR_dms[i])
vk_kpts[i] += weight * lib.dot(aoR_kband.T.conj(), tmp_Rq)
vkR_k1k2 = None
if dm_kpts.ndim == 3:
vj_kpts = vj_kpts[0]
vk_kpts = vk_kpts[0]
return lib.asarray(vj_kpts), lib.asarray(vk_kpts)
else:
for k2, kpt2 in enumerate(kpts):
aoR_dms = [lib.dot(aoR_kpts[k2], dms[i,k2]) for i in range(nset)]
for k1, kpt1 in enumerate(kpts):
vkR_k1k2 = get_vkR(mf, cell, aoR_kpts[k1], aoR_kpts[k2],
kpt1, kpt2)
for i in range(nset):
tmp_Rq = np.einsum('Rqs,Rs->Rq', vkR_k1k2, aoR_dms[i])
vk_kpts[i,k1] += weight * lib.dot(aoR_kpts[k1].T.conj(), tmp_Rq)
vkR_k1k2 = None
return vj_kpts.reshape(dm_kpts.shape), vk_kpts.reshape(dm_kpts.shape)
def get_ovlp(cell, grids=None):
if grids is None:
grids = gen_grid.BeckeGrids(cell)
grids.level = 3
grids.build()
aoR = numint.eval_ao(cell, grids.coords)
s = numpy.dot(aoR.T.conj(), grids.weights.reshape(-1, 1) * aoR).real
return s
def get_ovlp(cell, grids=None):
if grids is None:
grids = gen_grid.BeckeGrids(cell)
grids.level = 3
grids.build()
aoR = numint.eval_ao(cell, grids.coords)
s = numpy.dot(aoR.T.conj(), grids.weights.reshape(-1,1)*aoR).real
return s
def orbital(cell, outfile, coeff, nx=60, ny=60, nz=60):
'''Calculate orbital value on real space grid and write out in
CHGCAR format.
Args:
cell : Cell
pbc Cell
outfile : str
Name of Cube file to be written.
dm : ndarray
Density matrix of molecule.
Kwargs:
nx : int
Number of grid point divisions in x direction.
Note this is function of the molecule's size; a larger molecule
will have a coarser representation than a smaller one for the
same value.
ny : int
Number of grid point divisions in y direction.
nz : int
Number of grid point divisions in z direction.
Returns:
No return value. This function outputs a VASP chgcarlike file
(with phase if desired)...it can be opened in VESTA or VMD or
many other softwares
Examples:
>>> # generates the first MO from the list of mo_coefficents
>>> from pyscf.pbc import gto, scf
>>> from pyscf.tools import chgcar
>>> cell = gto.M(atom='H 0 0 0; H 0 0 1', a=numpy.eye(3)*3)
>>> mf = scf.RHF(cell).run()
>>> chgcar.orbital(cell, 'h2_mo1.CHGCAR', mf.mo_coeff[:,0])
'''
assert (isinstance(cell, pbcgto.Cell))
cc = CHGCAR(cell, nx=nx, ny=ny, nz=nz)
coords = cc.get_coords()
ngrids = cc.get_ngrids()
blksize = min(8000, ngrids)
orb_on_grid = numpy.empty(ngrids)
for ip0, ip1 in lib.prange(0, ngrids, blksize):
ao = numint.eval_ao(cell, coords[ip0:ip1])
orb_on_grid[ip0:ip1] = numpy.dot(ao, coeff)
orb_on_grid = orb_on_grid.reshape(nx, ny, nz)
cc.write(orb_on_grid,
outfile,
comment='Orbital value in real space (1/Bohr^3)')
def test_eval_mat(self):
cell, grids = make_grids(30)
ng = grids.weights.size
np.random.seed(1)
rho = np.random.random(ng)
rho *= 1/np.linalg.norm(rho)
vrho = np.random.random(ng)
ao1 = numint.eval_ao(cell, grids.coords)
mat1 = numint.eval_mat(cell, ao1, grids.weights, rho, vrho)
w = np.arange(mat1.size) * .01
self.assertAlmostEqual(np.dot(w,mat1.ravel()), (.14777107967912118+0j), 8)
def test_eval_mat(self):
cell, grids = make_grids([61]*3)
ng = grids.weights.size
np.random.seed(1)
rho = np.random.random(ng)
rho *= 1/np.linalg.norm(rho)
vrho = np.random.random(ng)
ao1 = numint.eval_ao(cell, grids.coords)
mat1 = numint.eval_mat(cell, ao1, grids.weights, rho, vrho)
w = np.arange(mat1.size) * .01
self.assertAlmostEqual(np.dot(w,mat1.ravel()), (.14777107967912118+0j), 8)
def get_pp_nl(cell, kpt=np.zeros(3)):
coords = gen_grid.gen_uniform_grids(cell)
aoR = numint.eval_ao(cell, coords, kpt)
nao = cell.nao_nr()
SI = cell.get_SI()
aokG = tools.fftk(np.asarray(aoR.T, order='C'),
cell.gs, np.exp(-1j*np.dot(coords, kpt))).T
ngs = len(aokG)
fakemol = pyscf.gto.Mole()
fakemol._atm = np.zeros((1,pyscf.gto.ATM_SLOTS), dtype=np.int32)
fakemol._bas = np.zeros((1,pyscf.gto.BAS_SLOTS), dtype=np.int32)
ptr = pyscf.gto.PTR_ENV_START
fakemol._env = np.zeros(ptr+10)
fakemol._bas[0,pyscf.gto.NPRIM_OF ] = 1
fakemol._bas[0,pyscf.gto.NCTR_OF ] = 1
fakemol._bas[0,pyscf.gto.PTR_EXP ] = ptr+3
fakemol._bas[0,pyscf.gto.PTR_COEFF] = ptr+4
Gv = np.asarray(cell.Gv+kpt)
G_rad = lib.norm(Gv, axis=1)
vppnl = np.zeros((nao,nao), dtype=np.complex128)
for ia in range(cell.natm):
symb = cell.atom_symbol(ia)
if symb not in cell._pseudo:
continue
pp = cell._pseudo[symb]
for l, proj in enumerate(pp[5:]):
rl, nl, hl = proj
if nl > 0:
hl = np.asarray(hl)
fakemol._bas[0,pyscf.gto.ANG_OF] = l
fakemol._env[ptr+3] = .5*rl**2
fakemol._env[ptr+4] = rl**(l+1.5)*np.pi**1.25
pYlm_part = pyscf.dft.numint.eval_ao(fakemol, Gv, deriv=0)
pYlm = np.empty((nl,l*2+1,ngs))
for k in range(nl):
qkl = pseudo.pp._qli(G_rad*rl, l, k)
pYlm[k] = pYlm_part.T * qkl
# pYlm is real
SPG_lmi = np.einsum('g,nmg->nmg', SI[ia].conj(), pYlm)
SPG_lm_aoG = np.einsum('nmg,gp->nmp', SPG_lmi, aokG)
tmp = np.einsum('ij,jmp->imp', hl, SPG_lm_aoG)
vppnl += np.einsum('imp,imq->pq', SPG_lm_aoG.conj(), tmp)
vppnl *= (1./ngs**2)
if aoR.dtype == np.double:
return vppnl.real
else:
return vppnl
def get_pp_nl(cell, kpt=np.zeros(3)):
coords = gen_grid.gen_uniform_grids(cell)
aoR = numint.eval_ao(cell, coords, kpt)
nao = cell.nao_nr()
SI = cell.get_SI()
aokG = tools.fftk(np.asarray(aoR.T, order='C'), cell.gs,
np.exp(-1j * np.dot(coords, kpt))).T
ngs = len(aokG)
fakemol = pyscf.gto.Mole()
fakemol._atm = np.zeros((1, pyscf.gto.ATM_SLOTS), dtype=np.int32)
fakemol._bas = np.zeros((1, pyscf.gto.BAS_SLOTS), dtype=np.int32)
ptr = pyscf.gto.PTR_ENV_START
fakemol._env = np.zeros(ptr + 10)
fakemol._bas[0, pyscf.gto.NPRIM_OF] = 1
fakemol._bas[0, pyscf.gto.NCTR_OF] = 1
fakemol._bas[0, pyscf.gto.PTR_EXP] = ptr + 3
fakemol._bas[0, pyscf.gto.PTR_COEFF] = ptr + 4
Gv = np.asarray(cell.Gv + kpt)
G_rad = lib.norm(Gv, axis=1)
vppnl = np.zeros((nao, nao), dtype=np.complex128)
for ia in range(cell.natm):
symb = cell.atom_symbol(ia)
if symb not in cell._pseudo:
continue
pp = cell._pseudo[symb]
for l, proj in enumerate(pp[5:]):
rl, nl, hl = proj
if nl > 0:
hl = np.asarray(hl)
fakemol._bas[0, pyscf.gto.ANG_OF] = l
fakemol._env[ptr + 3] = .5 * rl**2
fakemol._env[ptr + 4] = rl**(l + 1.5) * np.pi**1.25
pYlm_part = pyscf.dft.numint.eval_ao(fakemol, Gv, deriv=0)
pYlm = np.empty((nl, l * 2 + 1, ngs))
for k in range(nl):
qkl = pseudo.pp._qli(G_rad * rl, l, k)
pYlm[k] = pYlm_part.T * qkl
# pYlm is real
SPG_lmi = np.einsum('g,nmg->nmg', SI[ia].conj(), pYlm)
SPG_lm_aoG = np.einsum('nmg,gp->nmp', SPG_lmi, aokG)
tmp = np.einsum('ij,jmp->imp', hl, SPG_lm_aoG)
vppnl += np.einsum('imp,imq->pq', SPG_lm_aoG.conj(), tmp)
vppnl *= (1. / ngs**2)
if aoR.dtype == np.double:
return vppnl.real
else:
return vppnl
def get_nuc(cell, kpt=np.zeros(3)):
'''Get the bare periodic nuc-el AO matrix, with G=0 removed.
See Martin (12.16)-(12.21).
'''
coords = gen_grid.gen_uniform_grids(cell)
aoR = numint.eval_ao(cell, coords, kpt)
chargs = cell.atom_charges()
SI = cell.get_SI()
coulG = tools.get_coulG(cell)
vneG = -np.dot(chargs, SI) * coulG
vneR = tools.ifft(vneG, cell.gs).real
vne = np.dot(aoR.T.conj(), vneR.reshape(-1, 1) * aoR)
return vne
def get_nuc(cell, kpt=np.zeros(3)):
'''Get the bare periodic nuc-el AO matrix, with G=0 removed.
See Martin (12.16)-(12.21).
'''
coords = gen_grid.gen_uniform_grids(cell)
aoR = numint.eval_ao(cell, coords, kpt)
chargs = cell.atom_charges()
SI = cell.get_SI()
coulG = tools.get_coulG(cell)
vneG = -np.dot(chargs,SI) * coulG
vneR = tools.ifft(vneG, cell.gs).real
vne = np.dot(aoR.T.conj(), vneR.reshape(-1,1)*aoR)
return vne
def aoR_loop(self, cell, gs=None, kpts=None, kpt_band=None):
if kpts is None: kpts = self.kpts
kpts = numpy.asarray(kpts)
if gs is None:
gs = self.gs
else:
self.gs = gs
ngrids = numpy.prod(numpy.asarray(gs) * 2 + 1)
if (self._numint.cell is None or id(cell) != id(self._numint.cell)
or self._numint._deriv != 0
or self._numint._kpts.shape != kpts.shape
or abs(self._numint._kpts - kpts).sum() > 1e-9
or self._numint._coords.shape[0] != ngrids
or (gs is not None and numpy.any(gs != self.gs))):
nkpts = len(kpts)
coords = gen_grid.gen_uniform_grids(cell, gs)
nao = cell.nao_nr()
blksize = int(self.max_memory * 1e6 /
(nkpts * nao * 16 * numint.BLKSIZE)) * numint.BLKSIZE
blksize = min(max(blksize, numint.BLKSIZE), ngrids)
try:
self._numint.cache_ao(cell, kpts, 0, coords, nao, blksize)
except IOError as e:
sys.stderr.write(
'HDF5 file cannot be opened twice in different mode.\n')
raise e
with h5py.File(self._numint._ao.name, 'r') as f:
if kpt_band is None:
for k in range(len(kpts)):
aoR = f['ao/%d' % k].value
yield k, aoR
else:
kpt_band = numpy.reshape(kpt_band, 3)
where = numpy.argmin(pyscf.lib.norm(kpts - kpt_band, axis=1))
if abs(kpts[where] - kpt_band).sum() > 1e-9:
where = None
coords = gen_grid.gen_uniform_grids(cell, gs)
yield 0, numint.eval_ao(cell,
coords,
kpt_band,
deriv=deriv)
else:
yield where, f['ao/%d' % where].value
def aoR_loop(self, gs=None, kpts=None, kpt_band=None):
cell = self.cell
if kpts is None:
kpts = self.kpts
kpts = numpy.asarray(kpts)
if gs is None:
gs = self.gs
else:
self.gs = gs
ngrids = numpy.prod(numpy.asarray(gs) * 2 + 1)
if (
self._numint.cell is None
or id(cell) != id(self._numint.cell)
or self._numint._deriv != 0
or self._numint._kpts.shape != kpts.shape
or abs(self._numint._kpts - kpts).sum() > 1e-9
or self._numint._coords.shape[0] != ngrids
or (gs is not None and numpy.any(gs != self.gs))
):
nkpts = len(kpts)
coords = cell.gen_uniform_grids(gs)
nao = cell.nao_nr()
blksize = int(self.max_memory * 1e6 / (nkpts * nao * 16 * numint.BLKSIZE)) * numint.BLKSIZE
blksize = min(max(blksize, numint.BLKSIZE), ngrids)
try:
self._numint.cache_ao(cell, kpts, 0, coords, nao, blksize)
except IOError as e:
sys.stderr.write("HDF5 file cannot be opened twice in different mode.\n")
raise e
with h5py.File(self._numint._ao.name, "r") as f:
if kpt_band is None:
for k in range(len(kpts)):
aoR = f["ao/%d" % k].value
yield k, aoR
else:
kpt_band = numpy.reshape(kpt_band, 3)
where = numpy.argmin(lib.norm(kpts - kpt_band, axis=1))
if abs(kpts[where] - kpt_band).sum() > 1e-9:
where = None
coords = cell.gen_uniform_grids(gs)
yield 0, numint.eval_ao(cell, coords, kpt_band, deriv=0)
else:
yield where, f["ao/%d" % where].value
def get_pp(cell, kpt=np.zeros(3)):
'''Get the periodic pseudotential nuc-el AO matrix, with G=0 removed.
'''
coords = gen_grid.gen_uniform_grids(cell)
aoR = numint.eval_ao(cell, coords, kpt)
nao = cell.nao_nr()
SI = cell.get_SI()
vlocG = pseudo.get_vlocG(cell)
vlocG[:, 0] = 0
vpplocG = -np.sum(SI * vlocG, axis=0)
# vpploc evaluated in real-space
vpplocR = tools.ifft(vpplocG, cell.mesh)
vpploc = np.dot(aoR.T.conj(), vpplocR.reshape(-1, 1) * aoR)
# vppnonloc evaluated in reciprocal space
aokG = np.empty(aoR.shape, np.complex128)
expmikr = np.exp(-1j * np.dot(coords, kpt))
for i in range(nao):
aokG[:, i] = tools.fftk(aoR[:, i], cell.mesh, expmikr)
ngrids = len(aokG)
vppnl = np.zeros((nao, nao), dtype=np.complex128)
hs, projGs = pseudo.get_projG(cell, kpt)
for ia, [h_ia, projG_ia] in enumerate(zip(hs, projGs)):
for l, h in enumerate(h_ia):
nl = h.shape[0]
for m in range(-l, l + 1):
SPG_lm_aoG = np.zeros((nl, nao), dtype=np.complex128)
for i in range(nl):
SPG_lmi = SI[ia, :] * projG_ia[l][m][i]
SPG_lm_aoG[i, :] = np.einsum('g,gp->p', SPG_lmi.conj(),
aokG)
for i in range(nl):
for j in range(nl):
# Note: There is no (-1)^l here.
vppnl += h[i, j] * np.einsum('p,q->pq',
SPG_lm_aoG[i, :].conj(),
SPG_lm_aoG[j, :])
vppnl *= (1. / ngrids**2)
ovlp = cell.pbc_intor('int1e_ovlp_sph', hermi=1, kpts=kpt)
vpploc += 1. / cell.vol * np.sum(pseudo.get_alphas(cell)) * ovlp
return vpploc + vppnl
def get_pp(cell, kpt=np.zeros(3)):
'''Get the periodic pseudotential nuc-el AO matrix, with G=0 removed.
'''
coords = gen_grid.gen_uniform_grids(cell)
aoR = numint.eval_ao(cell, coords, kpt)
nao = cell.nao_nr()
SI = cell.get_SI()
vlocG = pseudo.get_vlocG(cell)
vlocG[:,0] = 0
vpplocG = -np.sum(SI * vlocG, axis=0)
# vpploc evaluated in real-space
vpplocR = tools.ifft(vpplocG, cell.mesh)
vpploc = np.dot(aoR.T.conj(), vpplocR.reshape(-1,1)*aoR)
# vppnonloc evaluated in reciprocal space
aokG = np.empty(aoR.shape, np.complex128)
expmikr = np.exp(-1j*np.dot(coords,kpt))
for i in range(nao):
aokG[:,i] = tools.fftk(aoR[:,i], cell.mesh, expmikr)
ngrids = len(aokG)
vppnl = np.zeros((nao,nao), dtype=np.complex128)
hs, projGs = pseudo.get_projG(cell, kpt)
for ia, [h_ia,projG_ia] in enumerate(zip(hs,projGs)):
for l, h in enumerate(h_ia):
nl = h.shape[0]
for m in range(-l,l+1):
SPG_lm_aoG = np.zeros((nl,nao), dtype=np.complex128)
for i in range(nl):
SPG_lmi = SI[ia,:] * projG_ia[l][m][i]
SPG_lm_aoG[i,:] = np.einsum('g,gp->p', SPG_lmi.conj(), aokG)
for i in range(nl):
for j in range(nl):
# Note: There is no (-1)^l here.
vppnl += h[i,j]*np.einsum('p,q->pq',
SPG_lm_aoG[i,:].conj(),
SPG_lm_aoG[j,:])
vppnl *= (1./ngrids**2)
ovlp = cell.pbc_intor('int1e_ovlp_sph', hermi=1, kpts=kpt)
vpploc += 1./cell.vol * np.sum(pseudo.get_alphas(cell)) * ovlp
return vpploc + vppnl
def export_unk(self, grid = [50,50,50]):
'''
Export the periodic part of BF in a real space grid for plotting with wannier90
'''
from scipy.io import FortranFile
grids_coor, weights = periodic_grid(self.cell, grid, order = 'F')
for k_id in range(self.num_kpts_loc):
spin = '.1'
if self.spin_up != None and self.spin_up == False : spin = '.2'
kpt = self.cell.get_abs_kpts(self.kpt_latt_loc[k_id])
ao = numint.eval_ao(self.cell, grids_coor, kpt = kpt)
u_ao = np.einsum('x,xi->xi', np.exp(-1j*np.dot(grids_coor, kpt)), ao, optimize = True)
unk_file = FortranFile('UNK' + "%05d" % (k_id + 1) + spin, 'w')
unk_file.write_record(np.asarray([grid[0], grid[1], grid[2], k_id + 1, self.num_bands_loc], dtype = np.int32))
mo_included = self.mo_coeff_kpts[k_id][:,self.band_included_list]
u_mo = np.einsum('xi,in->xn', u_ao, mo_included, optimize = True)
for band in range(len(self.band_included_list)):
unk_file.write_record(np.asarray(u_mo[:,band], dtype = np.complex128))
unk_file.close()
def get_A_mat(self):
'''
Construct the projection matrix: A_{m,n}^{\mathbf{k}}
Equation (62) in MV, Phys. Rev. B 56, 12847 or equation (22) in SMV, Phys. Rev. B 65, 035109
'''
A_matrix_loc = np.empty([self.num_kpts_loc, self.num_wann_loc, self.num_bands_loc], dtype = np.complex128)
if self.use_bloch_phases == True:
Amn = np.zeros([self.num_wann_loc, self.num_bands_loc])
np.fill_diagonal(Amn, 1)
A_matrix_loc[:,:,:] = Amn
else:
from pyscf.dft import numint,gen_grid
grids = gen_grid.Grids(self.cell).build()
coords = grids.coords
weights = grids.weights
for ith_wann in range(self.num_wann_loc):
frac_site = self.proj_site[ith_wann]
abs_site = frac_site.dot(self.real_lattice_loc) / param.BOHR
l = self.proj_l[ith_wann]
mr = self.proj_m[ith_wann]
r = self.proj_radial[ith_wann]
zona = self.proj_zona[ith_wann]
x_axis = self.proj_x[ith_wann]
z_axis = self.proj_z[ith_wann]
gr = g_r(coords, abs_site, l, mr, r, zona, x_axis, z_axis, unit = 'B')
ao_L0 = numint.eval_ao(self.cell, coords)
s_aoL0_g = np.einsum('i,i,iv->v', weights, gr, ao_L0, optimize = True)
for k_id in range(self.num_kpts_loc):
kpt = self.cell.get_abs_kpts(self.kpt_latt_loc[k_id])
mo_included = self.mo_coeff_kpts[k_id][:,self.band_included_list]
s_kpt = self.cell.pbc_intor('int1e_ovlp', hermi=1, kpts=kpt, pbcopt=lib.c_null_ptr())
s_ao = np.einsum('uv,v->u', s_kpt, s_aoL0_g, optimize = True)
A_matrix_loc[k_id,ith_wann,:] = np.einsum('v,vu,um->m', s_aoL0_g, s_kpt, mo_included, optimize = True).conj()
return A_matrix_loc
def get_pp_loc_part2(cell, kpt=np.zeros(3)):
coords = gen_grid.gen_uniform_grids(cell)
aoR = numint.eval_ao(cell, coords, kpt)
nao = cell.nao_nr()
SI = cell.get_SI()
G = lib.norm(cell.Gv, axis=1)
vlocG = np.zeros((cell.natm,len(G)))
for ia in range(cell.natm):
Zia = cell.atom_charge(ia)
symb = cell.atom_symbol(ia)
if symb not in cell._pseudo:
vlocG[ia] = 0
continue
pp = cell._pseudo[symb]
rloc, nexp, cexp = pp[1:3+1]
G_red = G*rloc
cfacs = np.array(
[1*G_red**0,
3 - G_red**2,
15 - 10*G_red**2 + G_red**4,
105 - 105*G_red**2 + 21*G_red**4 - G_red**6])
with np.errstate(divide='ignore'):
# Note the signs -- potential here is positive
vlocG[ia,:] = (# 4*np.pi * Zia * np.exp(-0.5*G_red**2)/G**2
- (2*np.pi)**(3/2.)*rloc**3*np.exp(-0.5*G_red**2)*(
np.dot(cexp, cfacs[:nexp])) )
vpplocG = -np.sum(SI * vlocG, axis=0)
vpplocR = tools.ifft(vpplocG, cell.gs).real
vpploc = np.dot(aoR.T.conj(), vpplocR.reshape(-1,1)*aoR)
if aoR.dtype == np.double:
return vpploc.real
else:
return vpploc
def get_ao_pairs_G(cell, kpt=np.zeros(3)):
'''Calculate forward (G|ij) and "inverse" (ij|G) FFT of all AO pairs.
Args:
cell : instance of :class:`Cell`
Returns:
ao_pairs_G, ao_pairs_invG : (ngrids, nao*(nao+1)/2) ndarray
The FFTs of the real-space AO pairs.
'''
coords = gen_uniform_grids(cell)
aoR = eval_ao(cell, coords, kpt) # shape = (coords, nao)
ngrids, nao = aoR.shape
gamma_point = abs(kpt).sum() < 1e-9
if gamma_point:
npair = nao * (nao + 1) // 2
ao_pairs_G = np.empty([ngrids, npair], np.complex128)
ij = 0
for i in range(nao):
for j in range(i + 1):
ao_ij_R = np.conj(aoR[:, i]) * aoR[:, j]
ao_pairs_G[:, ij] = tools.fft(ao_ij_R, cell.mesh)
#ao_pairs_invG[:,ij] = ngrids*tools.ifft(ao_ij_R, cell.mesh)
ij += 1
ao_pairs_invG = ao_pairs_G.conj()
else:
ao_pairs_G = np.zeros([ngrids, nao, nao], np.complex128)
for i in range(nao):
for j in range(nao):
ao_ij_R = np.conj(aoR[:, i]) * aoR[:, j]
ao_pairs_G[:, i, j] = tools.fft(ao_ij_R, cell.mesh)
ao_pairs_invG = ao_pairs_G.transpose(0, 2,
1).conj().reshape(-1, nao**2)
ao_pairs_G = ao_pairs_G.reshape(-1, nao**2)
return ao_pairs_G, ao_pairs_invG
def get_pp_loc_part2(cell, kpt=np.zeros(3)):
coords = gen_grid.gen_uniform_grids(cell)
aoR = numint.eval_ao(cell, coords, kpt)
nao = cell.nao_nr()
SI = cell.get_SI()
G = lib.norm(cell.Gv, axis=1)
vlocG = np.zeros((cell.natm,len(G)))
for ia in range(cell.natm):
Zia = cell.atom_charge(ia)
symb = cell.atom_symbol(ia)
if symb not in cell._pseudo:
vlocG[ia] = 0
continue
pp = cell._pseudo[symb]
rloc, nexp, cexp = pp[1:3+1]
G_red = G*rloc
cfacs = np.array(
[1*G_red**0,
3 - G_red**2,
15 - 10*G_red**2 + G_red**4,
105 - 105*G_red**2 + 21*G_red**4 - G_red**6])
with np.errstate(divide='ignore'):
# Note the signs -- potential here is positive
vlocG[ia,:] = (# 4*np.pi * Zia * np.exp(-0.5*G_red**2)/G**2
- (2*np.pi)**(3/2.)*rloc**3*np.exp(-0.5*G_red**2)*(
np.dot(cexp, cfacs[:nexp])) )
vpplocG = -np.sum(SI * vlocG, axis=0)
vpplocR = tools.ifft(vpplocG, cell.gs).real
vpploc = np.dot(aoR.T.conj(), vpplocR.reshape(-1,1)*aoR)
if aoR.dtype == np.double:
return vpploc.real
else:
return vpploc
def get_ao_pairs_G(cell, kpt=np.zeros(3)):
'''Calculate forward (G|ij) and "inverse" (ij|G) FFT of all AO pairs.
Args:
cell : instance of :class:`Cell`
Returns:
ao_pairs_G, ao_pairs_invG : (ngs, nao*(nao+1)/2) ndarray
The FFTs of the real-space AO pairs.
'''
coords = gen_uniform_grids(cell)
aoR = eval_ao(cell, coords, kpt) # shape = (coords, nao)
ngs, nao = aoR.shape
gamma_point = abs(kpt).sum() < 1e-9
if gamma_point:
npair = nao*(nao+1)//2
ao_pairs_G = np.empty([ngs, npair], np.complex128)
ij = 0
for i in range(nao):
for j in range(i+1):
ao_ij_R = np.conj(aoR[:,i]) * aoR[:,j]
ao_pairs_G[:,ij] = tools.fft(ao_ij_R, cell.gs)
#ao_pairs_invG[:,ij] = ngs*tools.ifft(ao_ij_R, cell.gs)
ij += 1
ao_pairs_invG = ao_pairs_G.conj()
else:
ao_pairs_G = np.zeros([ngs, nao,nao], np.complex128)
for i in range(nao):
for j in range(nao):
ao_ij_R = np.conj(aoR[:,i]) * aoR[:,j]
ao_pairs_G[:,i,j] = tools.fft(ao_ij_R, cell.gs)
ao_pairs_invG = ao_pairs_G.transpose(0,2,1).conj().reshape(-1,nao**2)
ao_pairs_G = ao_pairs_G.reshape(-1,nao**2)
return ao_pairs_G, ao_pairs_invG
else:
mf.kernel()
# end if
# grid density for molecular orbital
mydgs = 16
dgs = np.array([mydgs] * 3)
moR_fname = 'gs%d_' % mydgs + moR_fname
# run or load moR
if os.path.isfile(moR_fname):
moR = np.loadtxt(moR_fname)
else:
from pyscf.pbc.gto.cell import gen_uniform_grids
from pyscf.pbc.dft.numint import eval_ao
coords = gen_uniform_grids(cell, gs=dgs)
aoR = eval_ao(cell, coords)
moR = np.dot(aoR, mf.mo_coeff)
np.savetxt(moR_fname, moR)
# end if
mo_to_plot = [0, 1, 3, 4]
from qharv.inspect import volumetric
from skimage import measure
for iorb in mo_to_plot:
val = moR[:, iorb].reshape(2 * dgs + 1)
fval = volumetric.spline_volumetric(val)
grid = volumetric.axes_func_on_grid3d(axes, fval, grid_shape)
myup = default_up
mydn = default_dn
def kernel(mp, mo_energy=None, mo_coeff=None, verbose=logger.NOTE):
nao = mp.nmo
nocc = mp.nocc
nvirt = nao - nocc
ngs = mp.ngs
if mo_energy is None: mo_energy = mp.mo_energy #[nmo]
if mo_coeff is None: mo_coeff = mp.mo_coeff #[nao x nmo]
mo_energy = mo_energy.reshape(1, -1)
cell = mp._scf.cell
mesh = cell.mesh
smallmesh = mesh.copy()
smallmesh[-1] = int(numpy.floor(smallmesh[-1] / 2.0)) + 1
print "mesh: ", mesh
print "smallmesh: ", smallmesh
coulG = pbctools.get_coulG(cell, mesh=mesh) #[ngs]
coulG = coulG.reshape(mesh) #[mesh[0] x mesh[1] x mesh[2]]
coulG = coulG[:, :, :smallmesh[-1]].reshape([
numpy.product(smallmesh),
]) #[ngssmall]
coords = cell.gen_uniform_grids(mesh=mesh) #[ngs x 3]
(tauarray, weightarray, NLapPoints) = get_LT_data()
#batching
(max_grid_batch_size,
grid_batch_num, grid_batch) = get_batch(mp.max_memory,
ngs,
num_mat=3,
override=0) #batching grid
(max_mo_occ_batch_size,
mo_occ_batch_num, mo_occ_batch) = get_batch(mp.max_memory,
nocc,
num_mat=2,
override=0) #batching occ MOs
(max_mo_virt_batch_size, mo_virt_batch_num,
mo_virt_batch) = get_batch(mp.max_memory, nvirt, num_mat=2,
override=0) #batching virt MOs
(max_ao_batch_size, _, _) = get_batch(mp.max_memory,
nao,
num_mat=2,
override=0) #batching AOs
(shell_occ_batch,
ao_occ_batch) = get_ao_batch(cell.ao_loc_nr(),
max_ao_batch_size) #batching "occ" AOs
(shell_virt_batch,
ao_virt_batch) = get_ao_batch(cell.ao_loc_nr(),
max_ao_batch_size) #batching "virt" AOs
if grid_batch_num > 1:
print "Splitting the grid into ", grid_batch_num, " batches of max size ", max_grid_batch_size
if mo_occ_batch_num > 1:
print "Splitting the occupied MOs into ", mo_occ_batch_num, " batches of max size ", max_mo_occ_batch_size
if mo_virt_batch_num > 1:
print "Splitting the virtual MOs into ", mo_virt_batch_num, " batches of max size ", max_mo_virt_batch_size
Jtime = time.time()
EMP2J = 0.0
for i in range(numpy.amin([NLapPoints, mp.lt_points])):
Jint = 0.0
mo_occ = mo_coeff[:, :nocc] * numpy.exp(
-mo_energy[:, :nocc] * tauarray[i] / 2.) #[nao x nocc]
mo_virt = mo_coeff[:, nocc:] * numpy.exp(
mo_energy[:, nocc:] * tauarray[i] / 2.) #[nao x nvirt]
for b1 in range(grid_batch_num):
gbs = grid_batch[b1 + 1] - grid_batch[b1]
g_o = numpy.zeros((gbs, ngs), dtype='float64') #[gbs x ngs]
for a1 in range(mo_occ_batch_num):
mbs = mo_occ_batch[a1 + 1] - mo_occ_batch[a1]
moR_b = numpy.zeros((ngs, mbs), dtype='float64') #[ngs x mbs]
for a2 in range(len(shell_occ_batch) - 1):
aoR_b = numint.eval_ao(
cell,
coords,
shls_slice=(
shell_occ_batch[a2],
shell_occ_batch[a2 +
1])) #[ngs x shell_batch_size]
mo_occ_b = mo_occ[ao_occ_batch[a2]:ao_occ_batch[a2 + 1],
mo_occ_batch[a1]:mo_occ_batch[
a1 + 1]] #[shell_batch_size x mbs]
pylib.dot(aoR_b, mo_occ_b, alpha=1, c=moR_b,
beta=1) #[ngs x mbs]
aoR_b = mo_occ_b = None
pylib.dot(moR_b[grid_batch[b1]:grid_batch[b1 + 1]],
moR_b.T,
alpha=1,
c=g_o,
beta=1) #[gbs x ngs]
moR_b = None
F = numpy.zeros((gbs, ngs), dtype='float64') #[gbs x ngs]
for a1 in range(mo_virt_batch_num):
mbs = mo_virt_batch[a1 + 1] - mo_virt_batch[a1]
moR_b = numpy.zeros((ngs, mbs), dtype='float64') #[ngs x mbs]
for a2 in range(len(shell_virt_batch) - 1):
aoR_b = numint.eval_ao(
cell,
coords,
shls_slice=(
shell_virt_batch[a2],
shell_virt_batch[a2 +
1])) #[ngs x shell_batch_size]
mo_virt_b = mo_virt[ao_virt_batch[a2]:ao_virt_batch[a2 +
1],
mo_virt_batch[a1]:mo_virt_batch[
a1 + 1]] #[shell_batch_size x mbs]
pylib.dot(aoR_b, mo_virt_b, alpha=1, c=moR_b,
beta=1) #[ngs x mbs]
aoR_b = mo_virt_b = None
pylib.dot(moR_b[grid_batch[b1]:grid_batch[b1 + 1]],
moR_b.T,
alpha=1,
c=F,
beta=1) #[gbs x ngs]
moR_b = None
F *= g_o #[gbs x ngs]
g_o = None
if mp.optimization == "Cython":
fft_cython.getJ(gbs, F, coulG, mesh, smallmesh)
elif mp.optimization == "Python":
for j in range(gbs):
F[j] = compute_fft(F[j], coulG, mesh, smallmesh)
else:
raise RuntimeError('Only Cython and Python implemented!')
if grid_batch_num > 1:
for b2 in range(grid_batch_num):
if b2 != b1:
gbs_in = grid_batch[b2 + 1] - grid_batch[b2]
g_o = numpy.zeros((gbs_in, ngs),
dtype='float64') #[gbs_in x ngs]
for a1 in range(mo_occ_batch_num):
mbs = mo_occ_batch[a1 + 1] - mo_occ_batch[a1]
moR_b = numpy.zeros((ngs, mbs),
dtype='float64') #[ngs x mbs]
for a2 in range(len(shell_occ_batch) - 1):
aoR_b = numint.eval_ao(
cell,
coords,
shls_slice=(shell_occ_batch[a2],
shell_occ_batch[a2 + 1]
)) #[ngs x shell_batch_size]
mo_occ_b = mo_occ[
ao_occ_batch[a2]:ao_occ_batch[a2 + 1],
mo_occ_batch[a1]:mo_occ_batch[
a1 + 1]] #[shell_batch_size x mbs]
pylib.dot(aoR_b,
mo_occ_b,
alpha=1,
c=moR_b,
beta=1) #[ngs x mbs]
aoR_b = mo_occ_b = None
pylib.dot(moR_b[grid_batch[b2]:grid_batch[b2 + 1]],
moR_b.T,
alpha=1,
c=g_o,
beta=1) #[gbs_in x ngs]
moR_b = None
F_in = numpy.zeros((gbs_in, ngs),
dtype='float64') #[gbs_in x ngs]
for a1 in range(mo_virt_batch_num):
mbs = mo_virt_batch[a1 + 1] - mo_virt_batch[a1]
moR_b = numpy.zeros((ngs, mbs),
dtype='float64') #[ngs x mbs]
for a2 in range(len(shell_virt_batch) - 1):
aoR_b = numint.eval_ao(
cell,
coords,
shls_slice=(shell_virt_batch[a2],
shell_virt_batch[a2 + 1]
)) #[ngs x shell_batch_size]
mo_virt_b = mo_virt[
ao_virt_batch[a2]:ao_virt_batch[a2 + 1],
mo_virt_batch[a1]:mo_virt_batch[
a1 + 1]] #[shell_batch_size x mbs]
pylib.dot(aoR_b,
mo_virt_b,
alpha=1,
c=moR_b,
beta=1) #[ngs x mbs]
aoR_b = mo_virt_b = None
pylib.dot(moR_b[grid_batch[b2]:grid_batch[b2 + 1]],
moR_b.T,
alpha=1,
c=F_in,
beta=1) #[gbs_in x ngs]
moR_b = None
F_in *= g_o #[gbs_in x ngs]
g_o = None
if mp.optimization == "Cython":
fft_cython.getJ(gbs_in, F_in, coulG, mesh,
smallmesh)
elif mp.optimization == "Python":
for k in range(gbs_in):
F_in[k] = compute_fft(F_in[k], coulG, mesh,
smallmesh)
else:
raise RuntimeError(
'Only Cython and Python implemented!')
Jint += numpy.einsum(
'ij,ji->', F[:, grid_batch[b2]:grid_batch[b2 + 1]],
F_in[:, grid_batch[b1]:grid_batch[b1 + 1]])
F_in = None
else:
Jint += numpy.einsum(
'ij,ji->', F[:, grid_batch[b1]:grid_batch[b1 + 1]],
F[:, grid_batch[b1]:grid_batch[b1 + 1]])
F = None
else:
Jint += numpy.einsum('ij,ji->', F, F)
F = None
moRoccW = moRvirtW = None
EMP2J -= 2. * weightarray[i] * Jint * (cell.vol / ngs)**2.
print "Took this long for J: ", time.time() - Jtime
EMP2J = EMP2J.real
return EMP2J
def test_components(pseudo=None):
# The molecular calculation
mol = gto.Mole()
mol.unit = 'B'
L = 60
mol.atom.extend([
['He', (L / 2., L / 2., L / 2.)],
])
mol.basis = 'sto-3g'
mol.build()
m = rks.RKS(mol)
m.xc = 'LDA,VWN_RPA'
print(m.scf()) # -2.90705411168
dm = m.make_rdm1()
# The periodic calculation
cell = pbcgto.Cell()
cell.unit = 'B'
cell.a = np.diag([L, L, L])
cell.gs = np.array([80, 80, 80])
cell.atom = mol.atom
cell.basis = mol.basis
cell.pseudo = pseudo
cell.build()
# These should match reasonably well (roughly with accuracy of normalization)
print "Kinetic energy"
tao = pbchf.get_t(cell)
tao2 = mol.intor_symmetric('cint1e_kin_sph')
print np.dot(np.ravel(tao), np.ravel(dm)) # 2.82793077196
print np.dot(np.ravel(tao2), np.ravel(dm)) # 2.82352636524
print "Overlap"
sao = pbchf.get_ovlp(cell)
print np.dot(np.ravel(sao), np.ravel(dm)) # 1.99981725342
print np.dot(np.ravel(m.get_ovlp()), np.ravel(dm)) # 2.0
# The next two entries should *not* match, since G=0 component is removed
print "Coulomb (G!=0)"
jao = pbchf.get_j(cell, dm)
print np.dot(np.ravel(dm), np.ravel(jao)) # 4.03425518427
print np.dot(np.ravel(dm), np.ravel(m.get_j(dm))) # 4.22285177049
# The next two entries should *not* match, since G=0 component is removed
print "Nuc-el (G!=0)"
if cell.pseudo:
vppao = pbchf.get_pp(cell)
print np.dot(np.ravel(dm), np.ravel(vppao))
else:
neao = pbchf.get_nuc(cell)
print np.dot(np.ravel(dm), np.ravel(neao)) # -6.50203360062
vne = mol.intor_symmetric('cint1e_nuc_sph')
print np.dot(np.ravel(dm), np.ravel(vne)) # -6.68702326551
print "Normalization"
coords = gen_uniform_grids(cell)
aoR = eval_ao(cell, coords)
rhoR = eval_rho(cell, aoR, dm)
print cell.vol / len(rhoR) * np.sum(rhoR) # 1.99981725342 (should be 2.0)
print "(Hartree + vne) * DM"
print np.dot(np.ravel(dm), np.ravel(m.get_j(dm))) + np.dot(
np.ravel(dm), np.ravel(vne))
if cell.pseudo:
print np.einsum("ij,ij", dm, vppao + jao + pbchf.get_jvloc_G0(cell))
else:
print np.einsum("ij,ij", dm, neao + jao)
ew_cut = (40, 40, 40)
ew_eta = 0.05
for ew_eta in [0.1, 0.5, 1.]:
ew = pbchf.ewald(cell, ew_eta, ew_cut)
print "Ewald (eta, energy)", ew_eta, ew # should be same for all eta
print "Ewald divergent terms summation", ew
# These two should now match if the box is reasonably big to
# remove images, and ngs is big.
print "Total coulomb (analytic) ",
print(.5 * np.dot(np.ravel(dm), np.ravel(m.get_j(dm))) +
np.dot(np.ravel(dm), np.ravel(vne))) # -4.57559738004
if not cell.pseudo:
print "Total coulomb (fft coul + ewald)",
print np.einsum("ij,ij", dm, neao + .5 * jao) + ew # -4.57948259115
# Exc
cell.ew_eta, cell.ew_cut = ew_eta, ew_cut
mf = pbcdft.RKS(cell)
mf.xc = 'LDA,VWN_RPA'
rks.get_veff(mf, cell, dm)
print "Exc", mf._exc # -1.05967570089
def get_pp(cell, kpt=np.zeros(3)):
'''Get the periodic pseudotential nuc-el AO matrix
'''
import pyscf.dft
from pyscf.pbc import tools
from pyscf.pbc.dft import gen_grid
from pyscf.pbc.dft import numint
coords = gen_grid.gen_uniform_grids(cell)
aoR = numint.eval_ao(cell, coords, kpt)
nao = cell.nao_nr()
SI = cell.get_SI()
vlocG = get_vlocG(cell)
vpplocG = -np.sum(SI * vlocG, axis=0)
vpplocG[0] = np.sum(get_alphas(cell)) # from get_jvloc_G0 function
# vpploc evaluated in real-space
vpplocR = tools.ifft(vpplocG, cell.gs).real
vpploc = np.dot(aoR.T.conj(), vpplocR.reshape(-1, 1) * aoR)
# vppnonloc evaluated in reciprocal space
aokG = tools.fftk(np.asarray(aoR.T, order='C'), cell.gs,
np.exp(-1j * np.dot(coords, kpt))).T
ngs = len(aokG)
fakemol = pyscf.gto.Mole()
fakemol._atm = np.zeros((1, pyscf.gto.ATM_SLOTS), dtype=np.int32)
fakemol._bas = np.zeros((1, pyscf.gto.BAS_SLOTS), dtype=np.int32)
ptr = pyscf.gto.PTR_ENV_START
fakemol._env = np.zeros(ptr + 10)
fakemol._bas[0, pyscf.gto.NPRIM_OF] = 1
fakemol._bas[0, pyscf.gto.NCTR_OF] = 1
fakemol._bas[0, pyscf.gto.PTR_EXP] = ptr + 3
fakemol._bas[0, pyscf.gto.PTR_COEFF] = ptr + 4
Gv = np.asarray(cell.Gv + kpt)
G_rad = lib.norm(Gv, axis=1)
vppnl = np.zeros((nao, nao), dtype=np.complex128)
for ia in range(cell.natm):
symb = cell.atom_symbol(ia)
if symb not in cell._pseudo:
continue
pp = cell._pseudo[symb]
for l, proj in enumerate(pp[5:]):
rl, nl, hl = proj
if nl > 0:
hl = np.asarray(hl)
fakemol._bas[0, pyscf.gto.ANG_OF] = l
fakemol._env[ptr + 3] = .5 * rl**2
fakemol._env[ptr + 4] = rl**(l + 1.5) * np.pi**1.25
pYlm_part = pyscf.dft.numint.eval_ao(fakemol, Gv, deriv=0)
pYlm = np.empty((nl, l * 2 + 1, ngs))
for k in range(nl):
qkl = _qli(G_rad * rl, l, k)
pYlm[k] = pYlm_part.T * qkl
# pYlm is real
SPG_lmi = np.einsum('g,nmg->nmg', SI[ia].conj(), pYlm)
SPG_lm_aoG = np.einsum('nmg,gp->nmp', SPG_lmi, aokG)
tmp = np.einsum('ij,jmp->imp', hl, SPG_lm_aoG)
vppnl += np.einsum('imp,imq->pq', SPG_lm_aoG.conj(), tmp)
vppnl *= (1. / ngs**2)
if aoR.dtype == np.double:
return vpploc.real + vppnl.real
else:
return vpploc + vppnl
S_{k1,k2} = \int AO(k1,r) AO(k2,r) dr
'''
import numpy as np
from pyscf.pbc import gto
from pyscf.pbc import df
cell = gto.M(atom='H 0 0 0; H 1 0 1; H 0 1 0; H 1 .5 .5',
basis=[[0, (1, 1)], [1, (1.2, 1)]],
a=np.eye(3)*3)
k1, k2 = np.random.random((2,3))
#
# This is analytic FT to mimic this overlap integration
#
s_ft = df.ft_ao.ft_aopair(cell, -k1+k2, kpti_kptj=[k1,k2], q=np.zeros(3))
#
# Test the overlap using numerical integration
#
from pyscf.pbc.dft import gen_grid, numint
grids = gen_grid.UniformGrids(cell)
grids.build()
ao = numint.eval_ao(cell, grids.coords, kpt=k1)
u_k1 = np.einsum('x,xi->xi', np.exp(-1j*np.dot(grids.coords, k1)), ao)
ao = numint.eval_ao(cell, grids.coords, kpt=k2)
u_k2 = np.einsum('x,xi->xi', np.exp(-1j*np.dot(grids.coords, k2)), ao)
s_ref = np.einsum('x,xi,xj->ij', grids.weights, u_k1.conj(), u_k2)
print('max. Diff', abs(s_ref-s_ft).max())
def test_components(pseudo=None):
# The molecular calculation
mol = gto.Mole()
mol.unit = 'B'
L = 60
mol.atom.extend([['He', (L/2.,L/2.,L/2.)], ])
mol.basis = 'sto-3g'
mol.build()
m = rks.RKS(mol)
m.xc = 'LDA,VWN_RPA'
print(m.scf()) # -2.90705411168
dm = m.make_rdm1()
# The periodic calculation
cell = pbcgto.Cell()
cell.unit = 'B'
cell.a = np.diag([L,L,L])
cell.gs = np.array([80,80,80])
cell.atom = mol.atom
cell.basis = mol.basis
cell.pseudo = pseudo
cell.build()
# These should match reasonably well (roughly with accuracy of normalization)
print "Kinetic energy"
tao = pbchf.get_t(cell)
tao2 = mol.intor_symmetric('cint1e_kin_sph')
print np.dot(np.ravel(tao), np.ravel(dm)) # 2.82793077196
print np.dot(np.ravel(tao2), np.ravel(dm)) # 2.82352636524
print "Overlap"
sao = pbchf.get_ovlp(cell)
print np.dot(np.ravel(sao), np.ravel(dm)) # 1.99981725342
print np.dot(np.ravel(m.get_ovlp()), np.ravel(dm)) # 2.0
# The next two entries should *not* match, since G=0 component is removed
print "Coulomb (G!=0)"
jao = pbchf.get_j(cell, dm)
print np.dot(np.ravel(dm),np.ravel(jao)) # 4.03425518427
print np.dot(np.ravel(dm),np.ravel(m.get_j(dm))) # 4.22285177049
# The next two entries should *not* match, since G=0 component is removed
print "Nuc-el (G!=0)"
if cell.pseudo:
vppao = pbchf.get_pp(cell)
print np.dot(np.ravel(dm), np.ravel(vppao))
else:
neao = pbchf.get_nuc(cell)
print np.dot(np.ravel(dm), np.ravel(neao)) # -6.50203360062
vne = mol.intor_symmetric('cint1e_nuc_sph')
print np.dot(np.ravel(dm), np.ravel(vne)) # -6.68702326551
print "Normalization"
coords = gen_uniform_grids(cell)
aoR = eval_ao(cell, coords)
rhoR = eval_rho(cell, aoR, dm)
print cell.vol/len(rhoR)*np.sum(rhoR) # 1.99981725342 (should be 2.0)
print "(Hartree + vne) * DM"
print np.dot(np.ravel(dm),np.ravel(m.get_j(dm)))+np.dot(np.ravel(dm), np.ravel(vne))
if cell.pseudo:
print np.einsum("ij,ij", dm, vppao + jao + pbchf.get_jvloc_G0(cell))
else:
print np.einsum("ij,ij", dm, neao + jao)
ew_cut = (40,40,40)
ew_eta = 0.05
for ew_eta in [0.1, 0.5, 1.]:
ew = pbchf.ewald(cell, ew_eta, ew_cut)
print "Ewald (eta, energy)", ew_eta, ew # should be same for all eta
print "Ewald divergent terms summation", ew
# These two should now match if the box is reasonably big to
# remove images, and ngs is big.
print "Total coulomb (analytic) ",
print (.5*np.dot(np.ravel(dm), np.ravel(m.get_j(dm)))
+ np.dot(np.ravel(dm), np.ravel(vne)) ) # -4.57559738004
if not cell.pseudo:
print "Total coulomb (fft coul + ewald)",
print np.einsum("ij,ij", dm, neao + .5*jao) + ew # -4.57948259115
# Exc
cell.ew_eta, cell.ew_cut = ew_eta, ew_cut
mf = pbcdft.RKS(cell)
mf.xc = 'LDA,VWN_RPA'
rks.get_veff(mf, cell, dm)
print "Exc", mf._exc # -1.05967570089
def ao_on_grid(cell):
from pyscf.pbc.dft import gen_grid, numint
coords = gen_grid.gen_uniform_grids(cell)
aoR = numint.eval_ao(cell, coords)
return aoR
def get_pp(cell, kpt=np.zeros(3)):
'''Get the periodic pseudotential nuc-el AO matrix
'''
import pyscf.dft
from pyscf.pbc import tools
from pyscf.pbc.dft import gen_grid
from pyscf.pbc.dft import numint
coords = gen_grid.gen_uniform_grids(cell)
aoR = numint.eval_ao(cell, coords, kpt)
nao = cell.nao_nr()
SI = cell.get_SI()
vlocG = get_vlocG(cell)
vpplocG = -np.sum(SI * vlocG, axis=0)
vpplocG[0] = np.sum(get_alphas(cell)) # from get_jvloc_G0 function
# vpploc evaluated in real-space
vpplocR = tools.ifft(vpplocG, cell.gs).real
vpploc = np.dot(aoR.T.conj(), vpplocR.reshape(-1,1)*aoR)
# vppnonloc evaluated in reciprocal space
aokG = tools.fftk(np.asarray(aoR.T, order='C'),
cell.gs, np.exp(-1j*np.dot(coords, kpt))).T
ngs = len(aokG)
fakemol = pyscf.gto.Mole()
fakemol._atm = np.zeros((1,pyscf.gto.ATM_SLOTS), dtype=np.int32)
fakemol._bas = np.zeros((1,pyscf.gto.BAS_SLOTS), dtype=np.int32)
ptr = pyscf.gto.PTR_ENV_START
fakemol._env = np.zeros(ptr+10)
fakemol._bas[0,pyscf.gto.NPRIM_OF ] = 1
fakemol._bas[0,pyscf.gto.NCTR_OF ] = 1
fakemol._bas[0,pyscf.gto.PTR_EXP ] = ptr+3
fakemol._bas[0,pyscf.gto.PTR_COEFF] = ptr+4
Gv = np.asarray(cell.Gv+kpt)
G_rad = lib.norm(Gv, axis=1)
vppnl = np.zeros((nao,nao), dtype=np.complex128)
for ia in range(cell.natm):
symb = cell.atom_symbol(ia)
if symb not in cell._pseudo:
continue
pp = cell._pseudo[symb]
for l, proj in enumerate(pp[5:]):
rl, nl, hl = proj
if nl > 0:
hl = np.asarray(hl)
fakemol._bas[0,pyscf.gto.ANG_OF] = l
fakemol._env[ptr+3] = .5*rl**2
fakemol._env[ptr+4] = rl**(l+1.5)*np.pi**1.25
pYlm_part = pyscf.dft.numint.eval_ao(fakemol, Gv, deriv=0)
pYlm = np.empty((nl,l*2+1,ngs))
for k in range(nl):
qkl = _qli(G_rad*rl, l, k)
pYlm[k] = pYlm_part.T * qkl
# pYlm is real
SPG_lmi = np.einsum('g,nmg->nmg', SI[ia].conj(), pYlm)
SPG_lm_aoG = np.einsum('nmg,gp->nmp', SPG_lmi, aokG)
tmp = np.einsum('ij,jmp->imp', hl, SPG_lm_aoG)
vppnl += np.einsum('imp,imq->pq', SPG_lm_aoG.conj(), tmp)
vppnl *= (1./ngs**2)
if aoR.dtype == np.double:
return vpploc.real + vppnl.real
else:
return vpploc + vppnl
alat = 4.0
orbs_to_show = set((('2s', ''), ('2p', 'x'), ('3d', 'xy'), ('3d', 'z^2')))
assert len(orbs_to_show) == 4 # !!!! hard-code 4 plots
gs = np.array([mygs] * 3)
basis = bfd_basis()
cell = gto.M(a=alat * np.eye(3),
atom='C 2.0 2.0 2.0',
verbose=3,
gs=gs,
pseudo={'C': 'bfd'},
basis=basis)
coords = gen_grid.gen_uniform_grids(cell)
aoR = numint.eval_ao(cell, coords)
grid_shape = 2 * gs + 1
fig = plt.figure()
iplot = 0
bas_labels = cell.ao_labels(fmt=False)
for ibas in range(len(bas_labels)):
iatom, alabel, b_l, b_m = bas_labels[ibas]
if (b_l, b_m) not in orbs_to_show:
continue
iplot += 1
ax = fig.add_subplot(2, 2, iplot, projection='3d')
ax.set_title('%s %s %s' % (alabel, b_l, b_m))
