def test_kconserve(self):
cell = pbcgto.Cell()
cell.atom = 'He 0 0 0'
cell.a = '''0. 1.7834 1.7834
1.7834 0. 1.7834
1.7834 1.7834 0. '''
cell.build()
kpts = cell.make_kpts([3, 4, 5])
kconserve = tools.get_kconserv(cell, kpts)
self.assertAlmostEqual(lib.finger(kconserve), 84.88659638289468, 9)
def __init__(self, model, frozen=None):
"""
A mixin to support custom slices of mean-field attributes: `mo_coeff`, `mo_energy`, ...
PBC version.
Args:
model: the base model;
frozen (int, Iterable): the number of frozen valence orbitals or the list of frozen orbitals;
"""
self.model = model
self.space = format_frozen_k(frozen, len(model.mo_energy[0]), len(model.kpts))
self.kconserv = get_kconserv(self.model.cell, self.model.kpts).swapaxes(1, 2)
def build(self):
nmo_per_kpt = [x.size for x in self.o]
nocc_per_kpt = [np.sum(x) for x in self.o]
nmo, nocc = max(nmo_per_kpt), max(nocc_per_kpt)
nvir = nmo - nocc
self.opad, self.vpad = _padding_k_idx(nmo_per_kpt,
nocc_per_kpt,
kind='split')
self.kconserv = tools.get_kconserv(self.mf.cell, self.mf.kpts)
dtype = np.complex128 # TODO
self.eo = [e[o] for e, o in zip(self.e, self.o)]
self.ev = [e[v] for e, v in zip(self.e, self.v)]
self.co = [c[:, o] for c, o in zip(self.c, self.o)]
self.cv = [c[:, v] for c, v in zip(self.c, self.v)]
self.eo = np.array([e[x] for e, x in zip(self.eo, self.opad)])
self.ev = np.array([e[x] for e, x in zip(self.ev, self.vpad)])
self.co = np.array([c[:, x] for c, x in zip(self.co, self.opad)])
self.cv = np.array([c[:, x] for c, x in zip(self.cv, self.vpad)])
self.ovov = np.zeros((self.nkpts, ) * 3 + (nocc, nvir, nocc, nvir),
dtype=dtype)
self.ooov = np.zeros((self.nkpts, ) * 3 + (nocc, nocc, nocc, nvir),
dtype=dtype)
self.eija = np.zeros((self.nkpts, ) * 3 + (nocc, nocc, nvir))
self.t2 = self.ovov.copy()
for ki, kj, kk in mpi_helper.distr_iter(self.kpt_loop(3)):
kl = self.kconserv[ki, kj, kk]
kpts = [ki, kj, kk, kl]
eo, ev, co, cv = self.eo, self.ev, self.co, self.cv
self.ovov[ki, kj, kk] = self.ao2mo(kpts, co[ki], cv[kj], co[kk],
cv[kl])
self.ooov[ki, kj, kk] = self.ao2mo(kpts, co[ki], co[kj], co[kk],
cv[kl])
self.eija[ki, kj, kk] = lib.direct_sum('i,j,a->ija', eo[kj],
eo[kk], -ev[kl])
eiajb = lib.direct_sum('i,a,j,b->iajb', eo[ki], -ev[kj], eo[kk],
-ev[kl])
self.t2[ki, kj, kk] = self.ovov[ki, kj, kk] / eiajb
mpi_helper.barrier()
mpi_helper.allreduce_inplace(self.ovov)
mpi_helper.allreduce_inplace(self.ooov)
mpi_helper.allreduce_inplace(self.eija)
mpi_helper.allreduce_inplace(self.t2)
self._to_unpack = ['t2', 'ovov', 'ooov', 'eija']
def gen_eri_full():
kenec = dict()
kpts = mf.kpts
kconserv = tools.get_kconserv(cell, kpts)
for a, ka in enumerate(kpts):
for b, kb in enumerate(kpts):
for c, kc in enumerate(kpts):
d = kconserv[a, b, c]
kd = kpts[d]
eri_4d = mf.with_df.get_eri((ka, kb, kc, kd),
compact=False).reshape(
(nmo_kpts, ) * 4)
for i in range(nmo_kpts):
for j in range(nmo_kpts):
for k in range(nmo_kpts):
for l in range(nmo_kpts):
v = eri_4d[i, k, j, l]
#f.write('%s %s %s %s %s\n' % (a*nmo+i+1,b*nmo+j+1, c*nmo+k+1, d*nmo+l+1, trans(v)))
print((a * nmo_kpts + i, b * nmo_kpts + j,
c * nmo_kpts + k, d * nmo_kpts + l))
#kenec[(a*nmo_kpts+i,b*nmo_kpts+j, c*nmo_kpts+k, d*nmo_kpts+l) ] = trans(v)
kenec[(a * nmo_kpts + i, b * nmo_kpts + j,
c * nmo_kpts + k,
d * nmo_kpts + l)] = v
# 8 folds symmetrie
f = open('bielec_ao', 'w')
for l in range(nmo):
for k in range(nmo):
for j in range(l, nmo):
for i in range(max(j, k), nmo):
try:
v = kenec[(i, j, k, l)]
except KeyError:
pass
else:
if abs(v) > int_threshold:
#f.write('%s %s %s %s %s\n' % (i+1,j+1,k+1,l+1,trans(v)))
f.write('%s %s %s %s %s %s\n' %
(i + 1, j + 1, k + 1, l + 1, v.real,
v.imag))
def test_modified_cholesky(self):
cell = self.cell
mydf = df.FFTDF(cell, self.kpts)
comm = MPI.COMM_WORLD
nmo_tot = numpy.sum(self.nmo_pk)
maxvecs = 20 * nmo_tot
part = sc.Partition(comm, maxvecs, nmo_tot, self.nmo_max, self.nkpts)
gmap, Qi, ngs = sc.generate_grid_shifts(self.cell)
X = [numpy.identity(C.shape[0]) for C in self.mf.mo_coeff]
xik, xlj = sc.gen_orbital_products(cell, mydf, X, self.nmo_pk, ngs,
part, self.kpts, self.nmo_max)
kconserv = tools.get_kconserv(cell, self.kpts)
solver = sc.Cholesky(part, kconserv, 1e-3, verbose=False)
chol = solver.run(comm, xik, xlj, part, self.kpts, self.nmo_pk,
self.nmo_max, Qi, gmap)
self.assertEqual(chol.shape, (64, 64, 250))
self.assertAlmostEqual(numpy.max(numpy.abs(chol)),
0.807795378238119,
places=8)
self.assertEqual(len(chol[abs(chol) > 1e-10]), 101015)
C 0.8917 0.8917 0.8917
C 1.7834 1.7834 0.
C 2.6751 2.6751 0.8917
C 1.7834 0. 1.7834
C 2.6751 0.8917 2.6751
C 0. 1.7834 1.7834
C 0.8917 2.6751 2.6751''',
basis='gth-szv',
pseudo='gth-pade',
verbose=4,
)
nk = [2, 2, 2]
kpts = cell.make_kpts(nk)
kmf = scf.KRHF(cell, kpts, exxdiv=None)
kmf.kernel()
nmo = kmf.mo_coeff[0].shape[1]
kconserv = tools.get_kconserv(cell, kpts)
nkpts = len(kpts)
for kp in range(nkpts):
for kq in range(nkpts):
for kr in range(nkpts):
ks = kconserv[kp, kq, kr]
eri_kpt = kmf.with_df.ao2mo(
[kmf.mo_coeff[i] for i in (kp, kq, kr, ks)],
[kpts[i] for i in (kp, kq, kr, ks)])
eri_kpt = eri_kpt.reshape([nmo] * 4)
def write_hamil_supercell(comm,
scf_data,
hamil_file,
chol_cut,
verbose=True,
cas=None,
ortho_ao=False,
maxvecs=20,
exxdiv='ewald',
nelec=None):
nproc = comm.size
rank = comm.rank
tstart = time.clock()
# Unpack pyscf data.
# 1. core (1-body) Hamiltonian.
hcore = scf_data['hcore']
# 2. Rotation matrix to orthogonalised basis.
X = scf_data['X']
# 3. Pyscf cell object.
cell = scf_data['cell']
# 4. kpoints
kpts = scf_data['kpts']
# 5. MO occupancies.
Xocc = scf_data['Xocc']
# 6. Number of MOs per kpoint.
nmo_pk = scf_data['nmo_pk']
nkpts = len(kpts)
nao = cell.nao_nr()
# Update for GDF
mydf = df.FFTDF(cell, kpts)
nmo_max = numpy.max(nmo_pk)
assert (nmo_max is not None)
kconserv = tools.get_kconserv(cell, kpts)
ik2n, nmo_tot = setup_basis_map(Xocc, nmo_max, nkpts, nmo_pk, ortho_ao)
if comm.rank == 0:
h5file = h5py.File(hamil_file, "w")
h5grp = h5file.create_group("Hamiltonian")
else:
h5file = None
if rank == 0:
h1size = write_one_body(hcore, X, nkpts, nmo_pk, nmo_max, ik2n, h5grp)
comm.barrier()
numv = 0
if rank == 0 and verbose:
print(" # Time to reach cholesky: {:.2e} s".format(time.clock() -
tstart))
sys.stdout.flush()
tstart = time.clock()
# Setup parallel partition of work.
maxvecs = maxvecs * nmo_tot
part = Partition(comm, maxvecs, nmo_tot, nmo_max, nkpts)
if comm.rank == 0 and verbose:
print(" # Each kpoint is distributed accross {} mpi tasks.".format(
part.nproc_pk))
sys.stdout.flush()
maxvecs = part.maxvecs
# Set up mapping for shifted FFT grid.
gmap, Qi, ngs = generate_grid_shifts(cell)
if rank == 0 and verbose:
app_mem = (nkpts * nkpts * nmo_max * nmo_max * 2 * ngs *
16) / 1024.0**3
print(" # Approx. total memory required: {:.2e} GB.".format(app_mem))
sys.stdout.flush()
if rank == 0 and verbose:
print(" # Generating orbital products.")
sys.stdout.flush()
# Left and right pair densities. Only left contains Coulomb kernel.
t0 = time.clock()
Xaoik, Xaolj = gen_orbital_products(cell, mydf, X, nmo_pk, ngs, part, kpts,
nmo_max)
t1 = time.clock()
if part.rank == 0 and verbose:
print(" # Time to generate orbital products: {:.2e} s".format(t1 - t0))
sys.stdout.flush()
# Finally perform Cholesky decomposition.
solver = Cholesky(part, kconserv, gtol_chol=chol_cut)
cholvecs = solver.run(comm, Xaoik, Xaolj, part, kpts, nmo_pk, nmo_max, Qi,
gmap)
num_chol = cholvecs.shape[-1]
comm.barrier()
# Maximum number of integrals to include in single block in h5 file.
max_block_size = int(2e6)
# Write cholesky vectors to file.
v2cnts = write_cholesky(h5file, comm, cholvecs, part, ik2n, nmo_pk,
chol_cut, max_block_size)
if rank == 0:
# Write remaining descriptor arrays to file.
write_info(h5grp,
cell,
v2cnts,
h1size,
nmo_tot,
num_chol,
kpts,
ortho_ao,
Xocc,
nmo_pk,
ik2n,
exxdiv,
nelec=nelec)
h5file.close()
def pyscf2QP(cell, mf, kpts, kmesh=None, cas_idx=None, int_threshold=1E-8):
'''
kpts = List of kpoints coordinates. Cannot be null, for gamma is other script
kmesh = Mesh of kpoints (optional)
cas_idx = List of active MOs. If not specified all MOs are actives
int_threshold = The integral will be not printed in they are bellow that
'''
from pyscf.pbc import ao2mo
from pyscf.pbc import tools
from pyscf import lib
from pyscf.pbc.gto import ecp
mo_coef_threshold = int_threshold
ovlp_threshold = int_threshold
kin_threshold = int_threshold
ne_threshold = int_threshold
bielec_int_threshold = int_threshold
natom = len(cell.atom_coords())
print('n_atom', natom)
print('num_elec', cell.nelectron)
mo_coeff = mf.mo_coeff
# Mo_coeff actif
mo_k = np.array([c[:, cas_idx]
for c in mo_coeff] if cas_idx is not None else mo_coeff)
e_k = np.array(
[e[cas_idx]
for e in mf.mo_energy] if cas_idx is not None else mf.mo_energy)
Nk, nao, nmo = mo_k.shape
print("n Kpts", Nk)
print("n active Mos", nmo)
# Write all the parameter need to creat a dummy EZFIO folder who will containt the integral after.
# More an implentation detail than a real thing
with open('param', 'w') as f:
# Note the use of nmo_tot
f.write(' '.join(map(str,
(cell.nelectron * Nk, Nk * nmo, natom * Nk))))
with open('num_ao', 'w') as f:
f.write(str(nao * Nk))
# _
# |\ | _ | _ _. ._ |_) _ ._ | _ o _ ._
# | \| |_| (_ | (/_ (_| | | \ (/_ |_) |_| | _> | (_) | |
# |
#Total energy shift due to Ewald probe charge = -1/2 * Nelec*madelung/cell.vol =
shift = tools.pbc.madelung(cell, kpts) * cell.nelectron * -.5
e_nuc = (cell.energy_nuc() + shift) * Nk
print('nucl_repul', e_nuc)
with open('e_nuc', 'w') as f:
f.write(str(e_nuc))
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 mo_k2gamma(cell, mo_energy, mo_coeff, kpts, kmesh=None):
'''
Transform MOs in Kpoints to the equivalents supercell
'''
scell, phase = get_phase(cell, kpts, kmesh)
E_g = np.hstack(mo_energy)
C_k = np.asarray(mo_coeff)
Nk, Nao, Nmo = C_k.shape
NR = phase.shape[0]
# Transform AO indices
C_gamma = np.einsum('Rk, kum -> Rukm', phase, C_k)
C_gamma = C_gamma.reshape(Nao * NR, Nk * Nmo)
E_sort_idx = np.argsort(E_g)
E_g = E_g[E_sort_idx]
C_gamma = C_gamma[:, E_sort_idx]
s = scell.pbc_intor('int1e_ovlp')
assert (abs(
reduce(np.dot, (C_gamma.conj().T, s, C_gamma)) -
np.eye(Nmo * Nk)).max() < 1e-7)
# Transform MO indices
E_k_degen = abs(E_g[1:] - E_g[:-1]).max() < 1e-5
if np.any(E_k_degen):
degen_mask = np.append(False, E_k_degen) | np.append(
E_k_degen, False)
shift = min(E_g[degen_mask]) - .1
f = np.dot(C_gamma[:, degen_mask] * (E_g[degen_mask] - shift),
C_gamma[:, degen_mask].conj().T)
assert (abs(f.imag).max() < 1e-5)
e, na_orb = la.eigh(f.real, s, type=2)
C_gamma[:, degen_mask] = na_orb[:, e > 0]
if abs(C_gamma.imag).max() < 1e-7:
print('!Warning Some complexe pollutions in MOs are present')
C_gamma = C_gamma.real
if abs(
reduce(np.dot, (C_gamma.conj().T, s, C_gamma)) -
np.eye(Nmo * Nk)).max() < 1e-7:
print('!Warning Some complexe pollutions in MOs are present')
s_k = cell.pbc_intor('int1e_ovlp', kpts=kpts)
# overlap between k-point unitcell and gamma-point supercell
s_k_g = np.einsum('kuv,Rk->kuRv', s_k,
phase.conj()).reshape(Nk, Nao, NR * Nao)
# The unitary transformation from k-adapted orbitals to gamma-point orbitals
mo_phase = lib.einsum('kum,kuv,vi->kmi', C_k.conj(), s_k_g, C_gamma)
return mo_phase
# __ __ _
# |\/| | | | _ _ |_ _
# | | |__| |__ (_) (/_ | _>
#
with open('mo_coef_complex', 'w') as outfile:
c_kpts = np.reshape(mf.mo_coeff_kpts, (Nk, nao, nmo))
for ik in range(Nk):
shift1 = ik * nao + 1
shift2 = ik * nmo + 1
for i in range(nao):
for j in range(nmo):
cij = c_kpts[ik, i, j]
if abs(cij) > mo_coef_threshold:
outfile.write(
'%s %s %s %s\n' %
(i + shift1, j + shift2, cij.real, cij.imag))
# ___
# | ._ _|_ _ _ ._ _. | _ |\/| _ ._ _
# _|_ | | |_ (/_ (_| | (_| | _> | | (_) | | (_)
# _|
with open('overlap_ao_complex', 'w') as outfile:
#s_kpts_ao = np.reshape(mf.cell.pbc_intor('int1e_ovlp_sph',kpts=kpts),(Nk,nao,nao))
s_kpts_ao = np.reshape(mf.get_ovlp(cell=cell, kpts=kpts),
(Nk, nao, nao))
for ik in range(Nk):
shift = ik * nao + 1
for i in range(nao):
for j in range(i, nao):
sij = s_kpts_ao[ik, i, j]
if abs(sij) > ovlp_threshold:
outfile.write(
'%s %s %s %s\n' %
(i + shift, j + shift, sij.real, sij.imag))
with open('kinetic_ao_complex', 'w') as outfile:
#t_kpts_ao = np.reshape(mf.cell.pbc_intor('int1e_kin_sph',kpts=kpts),(Nk,nao,nao))
t_kpts_ao = np.reshape(cell.pbc_intor('int1e_kin', 1, 1, kpts=kpts),
(Nk, nao, nao))
for ik in range(Nk):
shift = ik * nao + 1
for i in range(nao):
for j in range(i, nao):
tij = t_kpts_ao[ik, i, j]
if abs(tij) > kin_threshold:
outfile.write(
'%s %s %s %s\n' %
(i + shift, j + shift, tij.real, tij.imag))
with open('ne_ao_complex', 'w') as outfile:
if mf.cell.pseudo:
v_kpts_ao = np.reshape(mf.with_df.get_pp(kpts=kpts),
(Nk, nao, nao))
else:
v_kpts_ao = np.reshape(mf.with_df.get_nuc(kpts=kpts),
(Nk, nao, nao))
if len(cell._ecpbas) > 0:
v_kpts_ao += np.reshape(ecp.ecp_int(cell, kpts), (Nk, nao, nao))
for ik in range(Nk):
shift = ik * nao + 1
for i in range(nao):
for j in range(i, nao):
vij = v_kpts_ao[ik, i, j]
if abs(vij) > ne_threshold:
outfile.write(
'%s %s %s %s\n' %
(i + shift, j + shift, vij.real, vij.imag))
# ___ _
# | ._ _|_ _ _ ._ _. | _ |_) o
# _|_ | | |_ (/_ (_| | (_| | _> |_) |
# _|
#
kconserv = tools.get_kconserv(cell, kpts)
def idx2_tri(i, j):
ij1 = min(i, j)
ij2 = max(i, j)
return ij1 + (ij2 * (ij2 - 1)) // 2
# eri_4d_ao = np.zeros((Nk,nao,Nk,nao,Nk,nao,Nk,nao), dtype=np.complex)
# for d, kd in enumerate(kpts):
# for c, kc in enumerate(kpts):
# if c > d: break
# idx2_cd = idx2_tri(c,d)
# for b, kb in enumerate(kpts):
# if b > d: break
# a = kconserv[b,c,d]
# if idx2_tri(a,b) > idx2_cd: continue
# if ((c==d) and (a>b)): continue
# ka = kpts[a]
# v = mf.with_df.get_ao_eri(kpts=[ka,kb,kc,kd],compact=False).reshape((nao,)*4)
# v *= 1./Nk
# eri_4d_ao[a,:,b,:,c,:,d] = v
#
# eri_4d_ao = eri_4d_ao.reshape([Nk*nao]*4)
#
# with open('bielec_ao_complex','w') as outfile:
# for d in range(Nk):
# for c in range(Nk):
# if c > d: break
# idx2_cd = idx2_tri(c,d)
# for b in range(Nk):
# if b > d: break
# a = kconserv[b,c,d]
# if idx2_tri(a,b) > idx2_cd: continue
# if ((c==d) and (a>b)): continue
# for l in range(nao):
# ll=l+d*nao
# for j in range(nao):
# jj=j+c*nao
# if jj>ll: break
# idx2_jjll = idx2_tri(jj,ll)
# for k in range(nao):
# kk=k+b*nao
# if kk>ll: break
# for i in range(nao):
# ii=k+a*nao
# if idx2_tri(ii,kk) > idx2_jjll: break
# if ((jj==ll) and (ii>kk)): break
# v=eri_4d_ao[ii,kk,jj,ll]
# if (abs(v) > bielec_int_threshold):
# outfile.write('%s %s %s %s %s %s\n' % (ii+1,jj+1,kk+1,ll+1,v.real,v.imag))
with open('bielec_ao_complex', 'w') as outfile:
for d, kd in enumerate(kpts):
for c, kc in enumerate(kpts):
if c > d: break
idx2_cd = idx2_tri(c, d)
for b, kb in enumerate(kpts):
if b > d: break
a = kconserv[b, c, d]
if idx2_tri(a, b) > idx2_cd: continue
if ((c == d) and (a > b)): continue
ka = kpts[a]
eri_4d_ao_kpt = mf.with_df.get_ao_eri(
kpts=[ka, kb, kc, kd], compact=False).reshape(
(nao, ) * 4)
eri_4d_ao_kpt *= 1. / Nk
for l in range(nao):
ll = l + d * nao
for j in range(nao):
jj = j + c * nao
if jj > ll: break
idx2_jjll = idx2_tri(jj, ll)
for k in range(nao):
kk = k + b * nao
if kk > ll: break
for i in range(nao):
ii = k + a * nao
if idx2_tri(ii, kk) > idx2_jjll: break
if ((jj == ll) and (ii > kk)): break
v = eri_4d_ao_kpt[i, k, j, l]
if (abs(v) > bielec_int_threshold):
outfile.write('%s %s %s %s %s %s\n' %
(ii + 1, jj + 1, kk + 1,
ll + 1, v.real, v.imag))
def pyscf2QP(cell,mf, kpts, kmesh=None, cas_idx=None, int_threshold = 1E-8,
print_ao_ints_bi=False,
print_mo_ints_bi=False,
print_ao_ints_df=True,
print_mo_ints_df=False,
print_ao_ints_mono=True,
print_mo_ints_mono=False):
'''
kpts = List of kpoints coordinates. Cannot be null, for gamma is other script
kmesh = Mesh of kpoints (optional)
cas_idx = List of active MOs. If not specified all MOs are actives
int_threshold = The integral will be not printed in they are bellow that
'''
from pyscf.pbc import ao2mo
from pyscf.pbc import tools
from pyscf.pbc.gto import ecp
import h5py
mo_coef_threshold = int_threshold
ovlp_threshold = int_threshold
kin_threshold = int_threshold
ne_threshold = int_threshold
bielec_int_threshold = int_threshold
natom = len(cell.atom_coords())
print('n_atom per kpt', natom)
print('num_elec per kpt', cell.nelectron)
mo_coeff = mf.mo_coeff
# Mo_coeff actif
mo_k = np.array([c[:,cas_idx] for c in mo_coeff] if cas_idx is not None else mo_coeff)
e_k = np.array([e[cas_idx] for e in mf.mo_energy] if cas_idx is not None else mf.mo_energy)
Nk, nao, nmo = mo_k.shape
print("n Kpts", Nk)
print("n active Mos per kpt", nmo)
print("n AOs per kpt", nao)
naux = mf.with_df.auxcell.nao
print("n df fitting functions", naux)
with open('num_df','w') as f:
f.write(str(naux))
# Write all the parameter need to creat a dummy EZFIO folder who will containt the integral after.
# More an implentation detail than a real thing
with open('param','w') as f:
# Note the use of nmo_tot
f.write(' '.join(map(str,(cell.nelectron*Nk, Nk*nmo, natom*Nk))))
with open('num_ao','w') as f:
f.write(str(nao*Nk))
with open('num_kpts','w') as f:
f.write(str(Nk))
# _
# |\ | _ | _ _. ._ |_) _ ._ | _ o _ ._
# | \| |_| (_ | (/_ (_| | | \ (/_ |_) |_| | _> | (_) | |
# |
#Total energy shift due to Ewald probe charge = -1/2 * Nelec*madelung/cell.vol =
shift = tools.pbc.madelung(cell, kpts)*cell.nelectron * -.5
e_nuc = (cell.energy_nuc() + shift)*Nk
print('nucl_repul', e_nuc)
with open('e_nuc','w') as f:
f.write(str(e_nuc))
# __ __ _
# |\/| | | | _ _ |_ _
# | | |__| |__ (_) (/_ | _>
#
with open('mo_coef_complex','w') as outfile:
c_kpts = np.reshape(mo_k,(Nk,nao,nmo))
for ik in range(Nk):
shift1=ik*nao+1
shift2=ik*nmo+1
for i in range(nao):
for j in range(nmo):
cij = c_kpts[ik,i,j]
if abs(cij) > mo_coef_threshold:
outfile.write('%s %s %s %s\n' % (i+shift1, j+shift2, cij.real, cij.imag))
# ___
# | ._ _|_ _ _ ._ _. | _ |\/| _ ._ _
# _|_ | | |_ (/_ (_| | (_| | _> | | (_) | | (_)
# _|
if mf.cell.pseudo:
v_kpts_ao = np.reshape(mf.with_df.get_pp(kpts=kpts),(Nk,nao,nao))
else:
v_kpts_ao = np.reshape(mf.with_df.get_nuc(kpts=kpts),(Nk,nao,nao))
if len(cell._ecpbas) > 0:
v_kpts_ao += np.reshape(ecp.ecp_int(cell, kpts),(Nk,nao,nao))
ne_ao = ('ne',v_kpts_ao,ne_threshold)
ovlp_ao = ('overlap',np.reshape(mf.get_ovlp(cell=cell,kpts=kpts),(Nk,nao,nao)),ovlp_threshold)
kin_ao = ('kinetic',np.reshape(cell.pbc_intor('int1e_kin',1,1,kpts=kpts),(Nk,nao,nao)),kin_threshold)
for name, intval_kpts_ao, thresh in (ne_ao, ovlp_ao, kin_ao):
if print_ao_ints_mono:
with open('%s_ao_complex' % name,'w') as outfile:
for ik in range(Nk):
shift=ik*nao+1
for i in range(nao):
for j in range(i,nao):
int_ij = intval_kpts_ao[ik,i,j]
if abs(int_ij) > thresh:
outfile.write('%s %s %s %s\n' % (i+shift, j+shift, int_ij.real, int_ij.imag))
if print_mo_ints_mono:
intval_kpts_mo = np.einsum('kim,kij,kjn->kmn',mo_k.conj(),intval_kpts_ao,mo_k)
with open('%s_mo_complex' % name,'w') as outfile:
for ik in range(Nk):
shift=ik*nmo+1
for i in range(nmo):
for j in range(i,nmo):
int_ij = intval_kpts_mo[ik,i,j]
if abs(int_ij) > thresh:
outfile.write('%s %s %s %s\n' % (i+shift, j+shift, int_ij.real, int_ij.imag))
# ___ _
# | ._ _|_ _ _ ._ _. | _ |_) o
# _|_ | | |_ (/_ (_| | (_| | _> |_) |
# _|
#
kconserv = tools.get_kconserv(cell, kpts)
with open('kconserv_complex','w') as outfile:
for a in range(Nk):
for b in range(Nk):
for c in range(Nk):
d = kconserv[a,b,c]
outfile.write('%s %s %s %s\n' % (a+1,c+1,b+1,d+1))
intfile=h5py.File(mf.with_df._cderi,'r')
j3c = intfile.get('j3c')
naosq = nao*nao
naotri = (nao*(nao+1))//2
j3ckeys = list(j3c.keys())
j3ckeys.sort(key=lambda strkey:int(strkey))
# in new(?) version of PySCF, there is an extra layer of groups before the datasets
# datasets used to be [/j3c/0, /j3c/1, /j3c/2, ...]
# datasets now are [/j3c/0/0, /j3c/1/0, /j3c/2/0, ...]
j3clist = [j3c.get(i+'/0') for i in j3ckeys]
if j3clist==[None]*len(j3clist):
# if using older version, stop before last level
j3clist = [j3c.get(i) for i in j3ckeys]
nkinvsq = 1./np.sqrt(Nk)
# dimensions are (kikj,iaux,jao,kao), where kikj is compound index of kpts i and j
# output dimensions should be reversed (nao, nao, naux, nkptpairs)
j3arr=np.array([(i.value.reshape([-1,nao,nao]) if (i.shape[1] == naosq) else makesq3(i.value,nao)) * nkinvsq for i in j3clist])
nkpt_pairs = j3arr.shape[0]
if print_ao_ints_df:
with open('df_ao_integral_array','w') as outfile:
pass
with open('df_ao_integral_array','a') as outfile:
for k,kpt_pair in enumerate(j3arr):
for iaux,dfbasfunc in enumerate(kpt_pair):
for i,i0 in enumerate(dfbasfunc):
for j,v in enumerate(i0):
if (abs(v) > bielec_int_threshold):
outfile.write('%s %s %s %s %s %s\n' % (i+1,j+1,iaux+1,k+1,v.real,v.imag))
if print_mo_ints_df:
kpair_list=[]
for i in range(Nk):
for j in range(Nk):
if(i>=j):
kpair_list.append((i,j,idx2_tri((i,j))))
j3mo = np.array([np.einsum('mij,ik,jl->mkl',j3arr[kij],mo_k[ki].conj(),mo_k[kj]) for ki,kj,kij in kpair_list])
with open('df_mo_integral_array','w') as outfile:
pass
with open('df_mo_integral_array','a') as outfile:
for k,kpt_pair in enumerate(j3mo):
for iaux,dfbasfunc in enumerate(kpt_pair):
for i,i0 in enumerate(dfbasfunc):
for j,v in enumerate(i0):
if (abs(v) > bielec_int_threshold):
outfile.write('%s %s %s %s %s %s\n' % (i+1,j+1,iaux+1,k+1,v.real,v.imag))
# eri_4d_ao = np.zeros((Nk,nao,Nk,nao,Nk,nao,Nk,nao), dtype=np.complex)
# for d, kd in enumerate(kpts):
# for c, kc in enumerate(kpts):
# if c > d: break
# idx2_cd = idx2_tri(c,d)
# for b, kb in enumerate(kpts):
# if b > d: break
# a = kconserv[b,c,d]
# if idx2_tri(a,b) > idx2_cd: continue
# if ((c==d) and (a>b)): continue
# ka = kpts[a]
# v = mf.with_df.get_ao_eri(kpts=[ka,kb,kc,kd],compact=False).reshape((nao,)*4)
# v *= 1./Nk
# eri_4d_ao[a,:,b,:,c,:,d] = v
#
# eri_4d_ao = eri_4d_ao.reshape([Nk*nao]*4)
if (print_ao_ints_bi or print_mo_ints_bi):
if print_ao_ints_bi:
with open('bielec_ao_complex','w') as outfile:
pass
if print_mo_ints_bi:
with open('bielec_mo_complex','w') as outfile:
pass
for d, kd in enumerate(kpts):
for c, kc in enumerate(kpts):
if c > d: break
idx2_cd = idx2_tri((c,d))
for b, kb in enumerate(kpts):
if b > d: break
a = kconserv[b,c,d]
if idx2_tri((a,b)) > idx2_cd: continue
if ((c==d) and (a>b)): continue
ka = kpts[a]
if print_ao_ints_bi:
with open('bielec_ao_complex','a') as outfile:
eri_4d_ao_kpt = mf.with_df.get_ao_eri(kpts=[ka,kb,kc,kd],compact=False).reshape((nao,)*4)
eri_4d_ao_kpt *= 1./Nk
for l in range(nao):
ll=l+d*nao
for j in range(nao):
jj=j+c*nao
if jj>ll: break
idx2_jjll = idx2_tri((jj,ll))
for k in range(nao):
kk=k+b*nao
if kk>ll: break
for i in range(nao):
ii=i+a*nao
if idx2_tri((ii,kk)) > idx2_jjll: break
if ((jj==ll) and (ii>kk)): break
v=eri_4d_ao_kpt[i,k,j,l]
if (abs(v) > bielec_int_threshold):
outfile.write('%s %s %s %s %s %s\n' % (ii+1,jj+1,kk+1,ll+1,v.real,v.imag))
if print_mo_ints_bi:
with open('bielec_mo_complex','a') as outfile:
eri_4d_mo_kpt = mf.with_df.ao2mo([mo_k[a], mo_k[b], mo_k[c], mo_k[d]],
[ka,kb,kc,kd],compact=False).reshape((nmo,)*4)
eri_4d_mo_kpt *= 1./Nk
for l in range(nmo):
ll=l+d*nmo
for j in range(nmo):
jj=j+c*nmo
if jj>ll: break
idx2_jjll = idx2_tri((jj,ll))
for k in range(nmo):
kk=k+b*nmo
if kk>ll: break
for i in range(nmo):
ii=i+a*nmo
if idx2_tri((ii,kk)) > idx2_jjll: break
if ((jj==ll) and (ii>kk)): break
v=eri_4d_mo_kpt[i,k,j,l]
if (abs(v) > bielec_int_threshold):
outfile.write('%s %s %s %s %s %s\n' % (ii+1,jj+1,kk+1,ll+1,v.real,v.imag))
def write_rhoG_supercell(comm, scf_data, hdf_file, Gcut,
verbose=True,
ortho_ao=False, write_real=False, phdf=False):
nproc = comm.size
rank = comm.rank
tstart = time.clock()
# Unpack pyscf data.
# 1. Rotation matrix to orthogonalised basis.
X = scf_data['X']
# 2. Pyscf cell object.
cell = scf_data['cell']
# 3. kpoints
kpts = scf_data['kpts']
# 4. Number of MOs per kpoint.
nmo_pk = scf_data['nmo_pk']
nkpts = len(kpts)
nao = cell.nao_nr()
quit()
# Update for GDF
mydf = df.FFTDF(cell, kpts)
nmo_max = numpy.max(nmo_pk)
assert(nmo_max is not None)
kconserv = tools.get_kconserv(cell, kpts)
ik2n, nmo_tot = setup_basis_map(Xocc, nmo_max, nkpts, nmo_pk, ortho_ao)
if comm.rank == 0:
h5file = h5py.File(hamil_file, "w")
h5grp = h5file.create_group("rhoG")
else:
h5file = None
numv = 0
# Setup parallel partition of work.
part = Partition(comm, 0, nmo_tot, nmo_max, nkpts)
if comm.rank == 0 and verbose:
print(" # Each kpoint is distributed accross {} mpi tasks."
.format(part.nproc_pk))
sys.stdout.flush()
# Set up mapping for shifted FFT grid.
if rank == 0 and verbose:
app_mem = (nmo_max*nmo_max*ngs*16) / 1024.0**3
print(" # Approx. local memory required: {:.2e} GB.".format(app_mem))
sys.stdout.flush()
if rank == 0 and verbose:
print(" # Generating orbital products.")
sys.stdout.flush()
for k in range(part.nkk):
k1 = part.n2k1[k]
k2 = part.n2k2[k]
i0 = part.ij0 // nmo_pk[k2]
iN = part.ijN // nmo_pk[k2]
if part.ijN % nmo_pk[k2] != 0:
iN += 1
if iN > nmo_pk[k1]:
iN = nmo_pk[k1]
pij = part.ij0 % nmo_pk[k2]
n_ = min(part.ijN, nmo_pk[k1]*nmo_pk[k2]) - part.ij0
X_t = X[k1][:,i0:iN].copy()
Xaoik[k,:,0:n_] = mydf.get_mo_pairs_G((X_t,X[k2]),
(kpts[k1],kpts[k2]),
kpts[k2]-kpts[k1],
compact=False)[:,pij:pij+n_]
X_t = None
Xaolj[k,:,:] = Xaoik[k,:,:]
coulG = tools.get_coulG(cell, kpts[k2]-kpts[k1], mesh=mydf.mesh)
Xaoik[k,:,:] *= (coulG*cell.vol/ngs**2).reshape(-1,1)
