def get_jk(mol_or_mf=None,
dm=None,
hermi=1,
with_j=True,
with_k=True,
omega=None):
'''MPI version of scf.hf.get_jk function'''
#vj = get_j(mol_or_mf, dm, hermi)
#vk = get_k(mol_or_mf, dm, hermi)
if isinstance(mol_or_mf, gto.mole.Mole):
mf = hf.SCF(mol_or_mf).view(SCF)
else:
mf = mol_or_mf
# dm may be too big for mpi4py library to serialize. Broadcast dm here.
if any(comm.allgather(dm is mpi.Message.SkippedArg)):
dm = mpi.bcast_tagged_array(dm)
mf.unpack_(comm.bcast(mf.pack()))
if mf.opt is None:
mf.opt = mf.init_direct_scf()
if omega is None:
vj, vk = _eval_jk(mf, dm, hermi, _jk_jobs_s8)
else:
with mf.mol.with_range_coulomb(omega):
vj, vk = _eval_jk(mf, dm, hermi, _jk_jobs_s8)
if rank == 0:
for i in range(vj.shape[0]):
lib.hermi_triu(vj[i], 1, inplace=True)
return vj.reshape(dm.shape), vk.reshape(dm.shape)
def get_k(mol_or_mf, dm, hermi=1):
if isinstance(mol_or_mf, gto.mole.Mole):
mf = hf.SCF(mol_or_mf).view(SCF)
else:
mf = mol_or_mf
# dm may be too big for mpi4py library to serialize. Broadcast dm here.
if any(comm.allgather(isinstance(dm, str) and dm == 'SKIPPED_ARG')):
dm = mpi.bcast_tagged_array(dm)
mf.unpack_(comm.bcast(mf.pack()))
if mf.opt is None:
mf.opt = mf.init_direct_scf()
with lib.temporary_env(mf.opt._this.contents,
fprescreen=_vhf._fpointer('CVHFnrs8_vk_prescreen')):
vk = _eval_jk(mf, dm, hermi, _vk_jobs_s8)
return vk.reshape(dm.shape)
def get_k(mol_or_mf=None, dm=None, hermi=1, omega=None):
if isinstance(mol_or_mf, gto.mole.Mole):
mf = hf.SCF(mol_or_mf).view(SCF)
else:
mf = mol_or_mf
# dm may be too big for mpi4py library to serialize. Broadcast dm here.
if any(comm.allgather(dm is mpi.Message.SkippedArg)):
dm = mpi.bcast_tagged_array(dm)
mf.unpack_(comm.bcast(mf.pack()))
if mf.opt is None:
mf.opt = mf.init_direct_scf()
with lib.temporary_env(mf.opt._this.contents,
fprescreen=_vhf._fpointer('CVHFnrs8_vk_prescreen')):
if omega is None:
vk = _eval_jk(mf, dm, hermi, _vk_jobs_s8)
else:
with mf.mol.with_range_coulomb(omega):
vk = _eval_jk(mf, dm, hermi, _vk_jobs_s8)
return vk.reshape(dm.shape)
def get_jk(mol_or_mf, dm, hermi=1):
'''MPI version of scf.hf.get_jk function'''
#vj = get_j(mol_or_mf, dm, hermi)
#vk = get_k(mol_or_mf, dm, hermi)
if isinstance(mol_or_mf, gto.mole.Mole):
mf = hf.SCF(mol_or_mf).view(SCF)
else:
mf = mol_or_mf
# dm may be too big for mpi4py library to serialize. Broadcast dm here.
if any(comm.allgather(isinstance(dm, str) and dm == 'SKIPPED_ARG')):
dm = mpi.bcast_tagged_array(dm)
mf.unpack_(comm.bcast(mf.pack()))
if mf.opt is None:
mf.opt = mf.init_direct_scf()
vj, vk = _eval_jk(mf, dm, hermi, _jk_jobs_s8)
if rank == 0:
for i in range(vj.shape[0]):
lib.hermi_triu(vj[i], 1, inplace=True)
return vj.reshape(dm.shape), vk.reshape(dm.shape)
def get_veff(mf, mol=None, dm=None, dm_last=0, vhf_last=0, hermi=1):
t0 = (time.clock(), time.time())
mf.unpack_(comm.bcast(mf.pack()))
mol = mf.mol
ni = mf._numint
if mf.nlc != '':
raise NotImplementedError
omega, alpha, hyb = ni.rsh_and_hybrid_coeff(mf.xc, spin=mol.spin)
if abs(omega) > 1e-10: # For range separated Coulomb operator
raise NotImplementedError
# Broadcast the large input arrays here.
if any(comm.allgather(isinstance(dm, str) and dm == 'SKIPPED_ARG')):
if rank == 0 and dm is None:
dm = mf.make_rdm1()
dm = mpi.bcast_tagged_array(dm)
if any(comm.allgather(isinstance(dm_last, str) and dm_last == 'SKIPPED_ARG')):
dm_last = mpi.bcast_tagged_array(dm_last)
if any(comm.allgather(isinstance(vhf_last, str) and vhf_last == 'SKIPPED_ARG')):
vhf_last = mpi.bcast_tagged_array(vhf_last)
ground_state = (isinstance(dm, numpy.ndarray) and dm.ndim == 2)
if mf.grids.coords is None:
_setup_grids_(mf, dm)
t0 = logger.timer(mf, 'setting up grids', *t0)
if hermi == 2: # because rho = 0
n, exc, vxc = 0, 0, 0
else:
n, exc, vxc = ni.nr_rks(mol, mf.grids, mf.xc, dm)
n = comm.allreduce(n)
exc = comm.allreduce(exc)
vxc = mpi.reduce(vxc)
logger.debug(mf, 'nelec by numeric integration = %s', n)
t0 = logger.timer(mf, 'vxc', *t0)
if abs(hyb) < 1e-10 and abs(alpha) < 1e-10:
vk = None
if getattr(vhf_last, 'vj', None) is not None:
ddm = numpy.asarray(dm) - dm_last
vj = mpi.reduce(mf.get_j(mol, ddm, hermi))
vj += vhf_last.vj
else:
vj = mf.get_j(mol, dm, hermi)
vj = mpi.reduce(vj)
vxc += vj
else:
if getattr(vhf_last, 'vk', None) is not None:
ddm = numpy.asarray(dm) - dm_last
vj, vk = mf.get_jk(mol, ddm, hermi)
ddm = None
vj = mpi.reduce(vj)
vk = mpi.reduce(vk) * hyb
vj += vhf_last.vj
vk += vhf_last.vk
else:
vj, vk = mf.get_jk(mol, dm, hermi)
vj = mpi.reduce(vj)
vk = mpi.reduce(vk) * hyb
vxc += vj - vk * .5
if ground_state:
exc -= numpy.einsum('ij,ji', dm, vk) * .5 * .5
if ground_state:
ecoul = numpy.einsum('ij,ji', dm, vj) * .5
else:
ecoul = None
vxc = lib.tag_array(vxc, ecoul=ecoul, exc=exc, vj=vj, vk=vk)
return vxc
def get_veff(mf, mol=None, dm=None, dm_last=0, vhf_last=0, hermi=1):
t0 = (logger.process_clock(), logger.perf_counter())
mf.unpack_(comm.bcast(mf.pack()))
mol = mf.mol
ni = mf._numint
if mf.nlc != '':
raise NotImplementedError
omega, alpha, hyb = ni.rsh_and_hybrid_coeff(mf.xc, spin=mol.spin)
# Broadcast the large input arrays here.
if any(comm.allgather(dm is mpi.Message.SkippedArg)):
if rank == 0 and dm is None:
dm = mf.make_rdm1()
dm = mpi.bcast_tagged_array(dm)
if any(comm.allgather(dm_last is mpi.Message.SkippedArg)):
dm_last = mpi.bcast_tagged_array(dm_last)
if any(comm.allgather(vhf_last is mpi.Message.SkippedArg)):
vhf_last = mpi.bcast_tagged_array(vhf_last)
ground_state = (dm.ndim == 3 and dm.shape[0] == 2)
if mf.grids.coords is None:
mpi_rks._setup_grids_(mf, dm[0]+dm[1])
t0 = logger.timer(mf, 'setting up grids', *t0)
if hermi == 2: # because rho = 0
n, exc, vxc = 0, 0, 0
else:
n, exc, vxc = ni.nr_uks(mol, mf.grids, mf.xc, dm)
n = comm.allreduce(n)
exc = comm.allreduce(exc)
vxc = mpi.reduce(vxc)
logger.debug(mf, 'nelec by numeric integration = %s', n)
t0 = logger.timer(mf, 'vxc', *t0)
if abs(hyb) < 1e-10 and abs(alpha) < 1e-10:
vk = None
if getattr(vhf_last, 'vj', None) is not None:
ddm = numpy.asarray(dm) - dm_last
ddm = ddm[0] + ddm[1]
vj = mf.get_j(mol, ddm, hermi)
vj += vhf_last.vj
else:
vj = mf.get_j(mol, dm[0]+dm[1], hermi)
vxc += vj
else:
if getattr(vhf_last, 'vk', None) is not None:
ddm = numpy.asarray(dm) - dm_last
vj, vk = mf.get_jk(mol, ddm, hermi)
vj = vj[0] + vj[1]
vk *= hyb
if abs(omega) > 1e-10:
vklr = mf.get_k(mol, ddm, hermi, omega=omega)
vk += vklr * (alpha - hyb)
ddm = None
vj += vhf_last.vj
vk += vhf_last.vk
else:
vj, vk = mf.get_jk(mol, dm, hermi)
vj = vj[0] + vj[1]
vk *= hyb
if abs(omega) > 1e-10:
vklr = mf.get_k(mol, dm, hermi, omega=omega)
vk += vklr * (alpha - hyb)
vxc += vj
vxc -= vk
if ground_state:
exc -=(numpy.einsum('ij,ji', dm[0], vk[0]) +
numpy.einsum('ij,ji', dm[1], vk[1])) * .5
if ground_state:
ecoul = numpy.einsum('ij,ji', dm[0]+dm[1], vj) * .5
else:
ecoul = None
vxc = lib.tag_array(vxc, ecoul=ecoul, exc=exc, vj=vj, vk=vk)
return vxc
def get_jk(mol_or_mf, dm, hermi, dmcur):
#, *args, **kwargs
'''MPI version of scf.hf.get_jk function'''
#vj = get_j(mol_or_mf, dm, hermi)
#vk = get_k(mol_or_mf, dm, hermi)
if isinstance(mol_or_mf, gto.mole.Mole):
mf = hf.SCF(mol_or_mf).view(SCF)
else:
mf = mol_or_mf
# dm may be too big for mpi4py library to serialize. Broadcast dm here.
if any(comm.allgather(isinstance(dm, str) and dm == 'SKIPPED_ARG')):
#dm = mpi.bcast_tagged_array_occdf(dm)
dm = mpi.bcast_tagged_array(dm)
mf.unpack_(comm.bcast(mf.pack()))
# initial and final grids level
grdlvl_i = 0
grdlvl_f = 1
# norm_ddm threshold for grids change
thrd_nddm = 0.01
# set block size to adapt memory
sblk = 100
# interspace betweeen v shell
intsp = 1
# threshold for u and v
gthrdu = 1e-7
gthrdvs = 1e-4
gthrdvd = 1e-4
global cond, wao_vx, ngridsx, coordsx, weightsx
dms = numpy.asarray(dm)
dm_shape = dms.shape
nao = dm_shape[-1]
dms = dms.reshape(-1,nao,nao)
nset = dms.shape[0]
vj = [0] * nset
vk = [0] * nset
# DF-J and sgX set
mf.with_df = mf
mol = mf.mol
global int2c, ovlp, ao_loc, rao_loc, ov_scr
# need set mf.initsgx = None in scf.SCF __init_
if mf.initsgx is None:
mf.initsgx = 0
cond = 0
# set auxbasis in input file, need self.auxbasis = None in __init__ of hf.py
# mf.auxbasis = 'weigend'
auxbasis = mf.auxbasis
auxbasis = comm.bcast(auxbasis)
mf.auxbasis = comm.bcast(mf.auxbasis)
auxmol = df.addons.make_auxmol(mol, auxbasis)
# (P|Q)
int2c = auxmol.intor('int2c2e', aosym='s1', comp=1)
if rank == 0: print('auxmol.basis',auxmol.basis,'number of aux basis',int2c.shape[0])
# for sgX
# ao_loc and rao_loc
intbn = mol._add_suffix('int3c2e')
intbn = gto.moleintor.ascint3(intbn)
ao_loc = gto.moleintor.make_loc(mol._bas, intbn)
rao_loc = numpy.zeros((nao),dtype=int)
ovlp = mol.intor_symmetric('int1e_ovlp')
for i in range(mol.nbas):
for jj in range(ao_loc[i],ao_loc[i+1]):
rao_loc[jj] = i
# ovlp screening
ov_scr = numpy.zeros((mol.nbas,mol.nbas),dtype=int)
for i in range(mol.nbas):
for j in range(mol.nbas):
if mol.bas_atom(i) == mol.bas_atom(j):
ov_scr[i,j] = 1
else:
movlp = numpy.amax(numpy.absolute(ovlp[ao_loc[i]:ao_loc[i+1],ao_loc[j]:ao_loc[j+1]]))
if movlp > gthrdvs:
ov_scr[i,j] = 1
if rank == 0: print('thrd_nddm',thrd_nddm, 'sblk',sblk, 'intsp',intsp, 'gthrdu',gthrdu)
if rank == 0: print('gthrdvs',gthrdvs, 'gthrdvd',gthrdvd)
# coase and fine grids change
grdchg = 0
norm_ddm = 0
for k in range(nset):
norm_ddm += numpy.linalg.norm(dms[k])
if norm_ddm < thrd_nddm and cond == 2 :
cond = 1
if cond == 0:
wao_vx, ngridsx, coordsx, weightsx = get_gridss(mol,grdlvl_i, sblk)
if rank == 0: print('grids level at first is', grdlvl_i)
cond = 2
elif cond == 1:
wao_vx, ngridsx, coordsx, weightsx = get_gridss(mol,grdlvl_f, sblk)
if rank == 0: print('grids level change to', grdlvl_f)
dms = numpy.asarray(dmcur)
dms = dms.reshape(-1,nao,nao)
grdchg = 1
cond = 3
if rank==0: Jtime=time.time()
# DF-J
dmtril = []
for k in range(nset):
dmtril.append(lib.pack_tril(dms[k]+dms[k].T))
i = numpy.arange(nao)
dmtril[k][i*(i+1)//2+i] *= .5
rho = []
b0 = 0
for eri1 in loop(mf.with_df):
naux, nao_pair = eri1.shape
#if rank==0: print('slice-naux',naux,'rank',rank)
b1 = b0 + naux
assert(nao_pair == nao*(nao+1)//2)
for k in range(nset):
if b0 == 0: rho.append(numpy.empty(paux[rank]))
rho[k][b0:b1] = numpy.dot(eri1, dmtril[k])
b0 = b1
orho = []
rec = []
for k in range(nset):
orho.append(mpi.gather(rho[k]))
if rank == 0:
ivj0 = scipy.linalg.solve(int2c, orho[k])
else:
ivj0 = None
rec.append(numpy.empty(paux[rank]))
comm.Scatterv([ivj0,paux],rec[k],root=0)
b0 = 0
for eri1 in loop(mf.with_df):
naux, nao_pair = eri1.shape
b1 = b0 + naux
assert(nao_pair == nao*(nao+1)//2)
for k in range(nset):
vj[k] += numpy.dot(rec[k][b0:b1].T, eri1)
b0 = b1
for k in range(nset):
vj[k] = comm.reduce(vj[k])
if rank==0: print "Took this long for J: ", time.time()-Jtime
if rank==0: Jtime=time.time()
# sgX
wao_v = wao_vx
coords = coordsx
weights = weightsx
for k in range(nset):
'''# Plus density screening
ov_scr = numpy.zeros((mol.nbas,mol.nbas),dtype=int)
for i in range(mol.nbas):
for j in range(mol.nbas):
movlp = numpy.amax(numpy.absolute(dms[k][ao_loc[i]:ao_loc[i+1],ao_loc[j]:ao_loc[j+1]]))
if movlp > gthrdvd:
ov_scr[i,j] = 1'''
# xfg screening
ngrids = coords.shape[0]
blksize = min(ngrids, sblk)
gscr = []
for i0, i1 in lib.prange(0, ngrids, blksize):
# screening u by value of grids
umaxg = numpy.amax(numpy.absolute(wao_v[i0:i1]), axis=0)
usi = numpy.argwhere(umaxg > gthrdu).reshape(-1)
if len(usi) != 0:
# screening v by ovlp then triangle matrix bn
uovl = ovlp[usi, :]
vmaxu = numpy.amax(numpy.absolute(uovl), axis=0)
osi = numpy.argwhere(vmaxu > gthrdvs).reshape(-1)
udms = dms[k][usi, :]
# screening v by dm and ovlp then triangle matrix bn
dmaxg = numpy.amax(numpy.absolute(udms), axis=0)
dsi = numpy.argwhere(dmaxg > gthrdvd).reshape(-1)
vsi = numpy.intersect1d(dsi, osi)
if len(vsi) != 0:
vsh = numpy.unique(rao_loc[vsi])
vshi = []
for i in range(vsh.shape[0]):
ista = ao_loc[vsh[i]]
iend = ao_loc[vsh[i]+1]
vshi.append(numpy.arange(ista, iend))
vshi = numpy.asarray(numpy.hstack(vshi))
dmsi = dms[k][usi]
fg = weights[i0:i1,None] * numpy.dot(wao_v[i0:i1,usi],dmsi[:,vshi])
gmaxfg = numpy.amax(numpy.absolute(fg), axis=1)
gmaxwao_v = numpy.amax(numpy.absolute(wao_v[i0:i1,usi]), axis=1)
gmaxtt = gmaxfg * gmaxwao_v
gscr0 = numpy.argwhere(gmaxtt > gthrdu).reshape(-1)
if gscr0.shape[0] > 0:
gscr.append(gscr0 + i0)
hgscr = numpy.hstack(gscr).reshape(-1)
coords = mpi.gather(coords[hgscr])
wao_v = mpi.gather(wao_v[hgscr])
weights = mpi.gather(weights[hgscr])
if rank == 0:
print('screened grids', coords.shape[0])
coords = numpy.array_split(coords, mpi.pool.size)
wao_v = numpy.array_split(wao_v, mpi.pool.size)
weights = numpy.array_split(weights, mpi.pool.size)
coords = mpi.scatter(coords)
wao_v = mpi.scatter(wao_v)
weights = mpi.scatter(weights)
# Kuv = Sum(Xug Avt Dkt Xkg)
ngrids = coords.shape[0]
for i0, i1 in lib.prange(0, ngrids, blksize):
# screening u by value of grids
umaxg = numpy.amax(numpy.absolute(wao_v[i0:i1]), axis=0)
usi = numpy.argwhere(umaxg > gthrdu).reshape(-1)
if len(usi) != 0:
# screening v by ovlp then triangle matrix bn
uovl = ovlp[usi, :]
vmaxu = numpy.amax(numpy.absolute(uovl), axis=0)
osi = numpy.argwhere(vmaxu > gthrdvs).reshape(-1)
udms = dms[k][usi, :]
# screening v by dm and ovlp then triangle matrix bn
dmaxg = numpy.amax(numpy.absolute(udms), axis=0)
dsi = numpy.argwhere(dmaxg > gthrdvd).reshape(-1)
vsi = numpy.intersect1d(dsi, osi)
if len(vsi) != 0:
vsh = numpy.unique(rao_loc[vsi])
nvsh = vsh.shape[0]
vov0 = ov_scr[vsh]
vov = vov0[:,vsh]
vshi = []
xvsh = vsh
ivx = [0]
vx = 0
for i in range(vsh.shape[0]):
ista = ao_loc[vsh[i]]
iend = ao_loc[vsh[i]+1]
vshi.append(numpy.arange(ista, iend))
vx += iend - ista
ivx.append(vx)
vshi = numpy.asarray(numpy.hstack(vshi))
nvshi = vshi.shape[0]
#print('ee',nvshi)
ivx = numpy.asarray(ivx)
vshbeg = vsh[0]
vshfin = vsh[-1]+1
dmsi = dms[k][usi]
fg = weights[i0:i1,None] * numpy.dot(wao_v[i0:i1,usi],dmsi[:,vshi])
fakemol = gto.fakemol_for_charges(coords[i0:i1])
#pmol = gto.mole.conc_mol(mol, fakemol)
intor = mol._add_suffix('int3c2e')
atm, bas, env = gto.mole.conc_env(mol._atm, mol._bas, mol._env,
fakemol._atm, fakemol._bas, fakemol._env)
shls_slice = (vshbeg, vshfin, vshbeg, vshfin, mol.nbas, mol.nbas+fakemol.nbas)
comp=1
#aosym='s1'
aosym='s2ij'
if aosym == 's2ij':
gv = getints3c_scr(intor, atm, bas, env, shls_slice, comp,
xvsh, nvshi, ivx, vov, fg, aosym)
else:
gv = getints3c_scr(intor, atm, bas, env, shls_slice, comp,
xvsh, nvshi, ivx, vov, fg, aosym)
vk0 = numpy.zeros((nao,nao))
vksp = lib.einsum('gu,gv->uv', wao_v[i0:i1,usi], gv)
vk1 = vk0[usi]
vk1[:,vshi] = vksp
vk0[usi] = vk1
vk[k] += vk0
wao_vw = weights[:,None] * wao_v
sn = lib.einsum('gu,gv->uv', wao_v, wao_vw)
vk[k] = comm.reduce(vk[k])
sn = comm.reduce(sn)
# SSn^-1 for grids to analitic
if rank == 0:
snsgk = scipy.linalg.solve(sn, vk[k])
vk[k] = numpy.matmul(ovlp, snsgk)
if hermi == 1:
vk[k] = (vk[k] + vk[k].T)*.5
if rank == 0:
print "Took this long for K: ", time.time()-Jtime
vj = lib.unpack_tril(numpy.asarray(vj), 1).reshape(dm_shape)
vk = numpy.asarray(vk).reshape(dm_shape)
return vj, vk, grdchg
