def comp_coulomb_pack(sv, ao_log=None, funct=coulomb_am, dtype=np.float64, **kw):
"""
Computes the matrix elements given by funct, for instance coulomb interaction
Args:
sv : (System Variables), this must have arrays of coordinates and species, etc
ao_log : description of functions (either orbitals or product basis functions)
Returns:
matrix elements for the whole system in packed form (lower triangular part)
"""
from pyscf.nao.m_ao_matelem import ao_matelem_c
from pyscf.nao.m_pack2den import cp_block_pack_u
aome = ao_matelem_c(sv.ao_log.rr, sv.ao_log.pp)
me = ao_matelem_c(sv.ao_log) if ao_log is None else aome.init_one_set(ao_log)
atom2s = np.zeros((sv.natm+1), dtype=np.int64)
for atom,sp in enumerate(sv.atom2sp): atom2s[atom+1]=atom2s[atom]+me.ao1.sp2norbs[sp]
norbs = atom2s[-1]
res = np.zeros(norbs*(norbs+1)//2, dtype=dtype)
for atom1,[sp1,rv1,s1,f1] in enumerate(zip(sv.atom2sp,sv.atom2coord,atom2s,atom2s[1:])):
#print("atom1 = {0}, rv1 = {1}".format(atom1, rv1))
for atom2,[sp2,rv2,s2,f2] in enumerate(zip(sv.atom2sp,sv.atom2coord,atom2s,atom2s[1:])):
if atom2>atom1: continue # skip
oo2f = funct(me,sp1,rv1,sp2,rv2, **kw)
cp_block_pack_u(oo2f, s1,f1,s2,f2, res)
return res, norbs
def overlap_lil(sv, ao_log=None, funct=overlap_ni,**kvargs):
"""
Computes the overlap matrix and returns it in List of Lists format (easy to index)
Args:
sv : (System Variables), this must have arrays of coordinates and species, etc
Returns:
overlap (real-space overlap) for the whole system
"""
from pyscf.nao.m_ao_matelem import ao_matelem_c
from scipy.sparse import lil_matrix
from numpy import array, int64, zeros
aome = ao_matelem_c(sv.ao_log.rr, sv.ao_log.pp)
me = aome.init_one_set(sv.ao_log) if ao_log is None else aome.init_one_set(ao_log)
atom2s = zeros((sv.natm+1), dtype=int64)
for atom,sp in enumerate(sv.atom2sp): atom2s[atom+1]=atom2s[atom]+me.ao1.sp2norbs[sp]
sp2rcut = array([max(mu2rcut) for mu2rcut in me.ao1.sp_mu2rcut])
n = atom2s[-1]
lil_mat = lil_matrix((n,n))
for atom1,[sp1,rv1,s1,f1] in enumerate(zip(sv.atom2sp,sv.atom2coord,atom2s,atom2s[1:])):
for atom2,[sp2,rv2,s2,f2] in enumerate(zip(sv.atom2sp,sv.atom2coord,atom2s,atom2s[1:])):
if (sp2rcut[sp1]+sp2rcut[sp2])**2<=sum((rv1-rv2)**2) : continue
lil_mat[s1:f1,s2:f2] = funct(me,sp1,rv1,sp2,rv2,**kvargs)
return lil_mat
def vxc_lil(sv, dm, xc_code, deriv, ao_log=None, dtype=float64, **kvargs):
"""
Computes the exchange-correlation matrix elements
Args:
sv : (System Variables), this must have arrays of coordinates and species, etc
Returns:
vxc,exc
"""
from pyscf.nao.m_xc_scalar_ni import xc_scalar_ni
from pyscf.nao.m_ao_matelem import ao_matelem_c
from scipy.sparse import lil_matrix
aome = ao_matelem_c(sv.ao_log.rr, sv.ao_log.pp, sv, dm)
me = aome.init_one_set(
sv.ao_log) if ao_log is None else aome.init_one_set(ao_log)
atom2s = zeros((sv.natm + 1), dtype=int64)
for atom, sp in enumerate(sv.atom2sp):
atom2s[atom + 1] = atom2s[atom] + me.ao1.sp2norbs[sp]
sp2rcut = array([max(mu2rcut) for mu2rcut in me.ao1.sp_mu2rcut])
lil = lil_matrix((atom2s[-1], atom2s[-1]), dtype=dtype)
for atom1, [sp1, rv1, s1, f1] in enumerate(
zip(sv.atom2sp, sv.atom2coord, atom2s, atom2s[1:])):
for atom2, [sp2, rv2, s2, f2] in enumerate(
zip(sv.atom2sp, sv.atom2coord, atom2s, atom2s[1:])):
if (sp2rcut[sp1] + sp2rcut[sp2])**2 <= sum((rv1 - rv2)**2):
continue
lil[s1:f1, s2:f2] = xc_scalar_ni(me, sp1, rv1, sp2, rv2, xc_code,
deriv, **kvargs)
return lil
def overlap_lil(sv, ao_log=None, funct=overlap_ni,**kvargs):
"""
Computes the overlap matrix and returns it in List of Lists format (easy to index)
Args:
sv : (System Variables), this must have arrays of coordinates and species, etc
Returns:
overlap (real-space overlap) for the whole system
"""
from pyscf.nao.m_ao_matelem import ao_matelem_c
from scipy.sparse import lil_matrix
from numpy import array, int64, zeros
aome = ao_matelem_c(sv.ao_log.rr, sv.ao_log.pp)
me = aome.init_one_set(sv.ao_log) if ao_log is None else aome.init_one_set(ao_log)
atom2s = zeros((sv.natm+1), dtype=int64)
for atom,sp in enumerate(sv.atom2sp): atom2s[atom+1]=atom2s[atom]+me.ao1.sp2norbs[sp]
sp2rcut = array([max(mu2rcut) for mu2rcut in me.ao1.sp_mu2rcut])
n = atom2s[-1]
lil_mat = lil_matrix((n,n))
for atom1,[sp1,rv1,s1,f1] in enumerate(zip(sv.atom2sp,sv.atom2coord,atom2s,atom2s[1:])):
for atom2,[sp2,rv2,s2,f2] in enumerate(zip(sv.atom2sp,sv.atom2coord,atom2s,atom2s[1:])):
if (sp2rcut[sp1]+sp2rcut[sp2])**2<=sum((rv1-rv2)**2) : continue
lil_mat[s1:f1,s2:f2] = funct(me,sp1,rv1,sp2,rv2,**kvargs)
return lil_mat
def comp_coulomb_den(sv, ao_log=None, funct=coulomb_am, dtype=np.float64, **kvargs):
"""
Computes the matrix elements given by funct, for instance coulomb interaction
Args:
sv : (System Variables), this must have arrays of coordinates and species, etc
Returns:
matrix elements (real-space overlap) for the whole system
"""
from pyscf.nao.m_ao_matelem import ao_matelem_c
aome = ao_matelem_c(sv.ao_log.rr, sv.ao_log.pp)
me = aome.init_one_set(sv.ao_log) if ao_log is None else aome.init_one_set(ao_log)
atom2s = np.zeros((sv.natm+1), dtype=np.int32)
for atom,sp in enumerate(sv.atom2sp): atom2s[atom+1]=atom2s[atom]+me.ao1.sp2norbs[sp]
norbs = atom2s[-1]
# dim triangular matrix: n*(n+1)/2
res = np.zeros((norbs,norbs), dtype=dtype)
for atom1,[sp1,rv1,s1,f1] in enumerate(zip(sv.atom2sp,sv.atom2coord,atom2s,atom2s[1:])):
for atom2,[sp2,rv2,s2,f2] in enumerate(zip(sv.atom2sp,sv.atom2coord,atom2s,atom2s[1:])):
oo2f = funct(me,sp1,rv1,sp2,rv2,**kvargs)
res[s1:f1,s2:f2] = oo2f[:,:]
#np.savetxt("kernel_pyscf_dens.txt", res)
#print("sum(kernel) = ", np.sum(abs(res)))
#import sys
#sys.exit()
return res
def comp_coulomb_den(sv,
ao_log=None,
funct=coulomb_am,
dtype=np.float64,
**kvargs):
"""
Computes the matrix elements given by funct, for instance coulomb interaction
Args:
sv : (System Variables), this must have arrays of coordinates and species, etc
Returns:
matrix elements (real-space overlap) for the whole system
"""
from pyscf.nao.m_ao_matelem import ao_matelem_c
aome = ao_matelem_c(sv.ao_log.rr, sv.ao_log.pp)
me = aome.init_one_set(
sv.ao_log) if ao_log is None else aome.init_one_set(ao_log)
atom2s = np.zeros((sv.natm + 1), dtype=np.int32)
for atom, sp in enumerate(sv.atom2sp):
atom2s[atom + 1] = atom2s[atom] + me.ao1.sp2norbs[sp]
norbs = atom2s[-1]
# dim triangular matrix: n*(n+1)/2
res = np.zeros((norbs, norbs), dtype=dtype)
for atom1, [sp1, rv1, s1, f1] in enumerate(
zip(sv.atom2sp, sv.atom2coord, atom2s, atom2s[1:])):
for atom2, [sp2, rv2, s2, f2] in enumerate(
zip(sv.atom2sp, sv.atom2coord, atom2s, atom2s[1:])):
oo2f = funct(me, sp1, rv1, sp2, rv2, **kvargs)
res[s1:f1, s2:f2] = oo2f[:, :]
return res
def comp_coulomb_pack(sv,
ao_log=None,
funct=coulomb_am,
dtype=np.float64,
**kvargs):
"""
Computes the matrix elements given by funct, for instance coulomb interaction
Args:
sv : (System Variables), this must have arrays of coordinates and species, etc
ao_log : description of functions (either orbitals or product basis functions)
Returns:
matrix elements (real-space overlap) for the whole system
"""
from pyscf.nao.m_ao_matelem import ao_matelem_c
from pyscf.nao.m_pack2den import triu_indices
aome = ao_matelem_c(sv.ao_log.rr, sv.ao_log.pp)
me = ao_matelem_c(
sv.ao_log) if ao_log is None else aome.init_one_set(ao_log)
atom2s = np.zeros((sv.natm + 1), dtype=np.int64)
for atom, sp in enumerate(sv.atom2sp):
atom2s[atom + 1] = atom2s[atom] + me.ao1.sp2norbs[sp]
norbs = atom2s[-1]
res = np.zeros(norbs * (norbs + 1) // 2, dtype=dtype)
ind = triu_indices(norbs)
for atom1, [sp1, rv1, s1, f1] in enumerate(
zip(sv.atom2sp, sv.atom2coord, atom2s, atom2s[1:])):
for atom2, [sp2, rv2, s2, f2] in enumerate(
zip(sv.atom2sp, sv.atom2coord, atom2s, atom2s[1:])):
#if atom2>atom1: continue # skip
oo2f = funct(me, sp1, rv1, sp2, rv2, **kvargs)
if use_numba:
fill_triu(oo2f, ind, res, s1, f1, s2, f2)
else:
for i1 in range(s1, f1):
for i2 in range(s2, f2):
if ind[i1, i2] >= 0:
res[ind[i1, i2]] = oo2f[i1 - s1, i2 - s2]
return res, norbs
def vxc_pack(self, **kw):
"""
Computes the exchange-correlation matrix elements packed version (upper triangular)
Args:
sv : (System Variables), this must have arrays of coordinates and species, etc
Returns:
vxc,exc
"""
from pyscf.nao.m_xc_scalar_ni import xc_scalar_ni
from pyscf.nao.m_ao_matelem import ao_matelem_c
#sv, dm, xc_code, deriv, kernel=None, ao_log=None, dtype=float64, **kvargs
sv = self
dm = kw['dm'] if 'dm' in kw else self.make_rdm1()
kernel = kw['kernel'] if 'kernel' in kw else None
ao_log = kw['ao_log'] if 'ao_log' in kw else self.ao_log
(xc_code,iskw) = (kw['xc_code'],True) if 'xc_code' in kw else (self.xc_code,False)
dtype = kw['dtype'] if 'dtype' in kw else float64
aome = ao_matelem_c(ao_log.rr, ao_log.pp, sv, dm)
me = aome.init_one_set(ao_log)
atom2s = zeros((sv.natm+1), dtype=int64)
for atom,sp in enumerate(sv.atom2sp): atom2s[atom+1]=atom2s[atom]+me.ao1.sp2norbs[sp]
sp2rcut = array([max(mu2rcut) for mu2rcut in me.ao1.sp_mu2rcut])
norbs = atom2s[-1]
#ind = triu_indices(norbs)
if kernel is None: kernel = zeros(norbs*(norbs+1)//2, dtype=dtype)
if kernel.size != int(norbs*(norbs+1)//2):
print('kernel.size ', kernel.size)
print('norbs*(norbs+1)//2 ', int(norbs*(norbs+1)//2))
raise ValueError("wrong dimension for kernel")
for atom1,[sp1,rv1,s1,f1] in enumerate(zip(sv.atom2sp,sv.atom2coord,atom2s,atom2s[1:])):
for atom2,[sp2,rv2,s2,f2] in enumerate(zip(sv.atom2sp,sv.atom2coord,atom2s,atom2s[1:])):
if atom2>atom1: continue
if (sp2rcut[sp1]+sp2rcut[sp2])**2<=sum((rv1-rv2)**2) : continue
xc = xc_scalar_ni(me,sp1,rv1,sp2,rv2,**kw) if iskw else xc_scalar_ni(me,sp1,rv1,sp2,rv2,xc_code=xc_code,**kw)
if use_numba:
fill_triu_v2(xc, kernel, s1, f1, s2, f2, norbs, add=True)
else:
for i1 in range(s1,f1):
for i2 in range(s2, min(i1+1, f2)):
ind = 0
if i2 > 0:
for beta in range(1, i2+1):
ind += norbs -beta
ind += i1
kernel[ind] += xc[i1-s1,i2-s2]
return kernel
def comp_coulomb_pack(sv,
ao_log=None,
funct=coulomb_am,
dtype=np.float64,
**kvargs):
"""
Computes the matrix elements given by funct, for instance coulomb interaction
Args:
sv : (System Variables), this must have arrays of coordinates and species, etc
ao_log : description of functions (either orbitals or product basis functions)
Returns:
matrix elements for the whole system in packed form (lower triangular part)
"""
from pyscf.nao.m_ao_matelem import ao_matelem_c
from pyscf.nao.m_pack2den import ij2pack_l
aome = ao_matelem_c(sv.ao_log.rr, sv.ao_log.pp)
me = ao_matelem_c(
sv.ao_log) if ao_log is None else aome.init_one_set(ao_log)
atom2s = np.zeros((sv.natm + 1), dtype=np.int64)
for atom, sp in enumerate(sv.atom2sp):
atom2s[atom + 1] = atom2s[atom] + me.ao1.sp2norbs[sp]
norbs = atom2s[-1]
res = np.zeros(norbs * (norbs + 1) // 2, dtype=dtype)
for atom1, [sp1, rv1, s1, f1] in enumerate(
zip(sv.atom2sp, sv.atom2coord, atom2s, atom2s[1:])):
for atom2, [sp2, rv2, s2, f2] in enumerate(
zip(sv.atom2sp, sv.atom2coord, atom2s, atom2s[1:])):
if atom2 > atom1: continue # skip
oo2f = funct(me, sp1, rv1, sp2, rv2, **kvargs)
if use_numba:
fill_triu_v2(oo2f, res, s1, f1, s2, f2, norbs)
else:
for i1 in range(s1, f1):
for i2 in range(s2, min(i1 + 1, f2)):
res[ij2pack_l(i1, i2, norbs)] = oo2f[i1 - s1, i2 - s2]
#print("sum kernel: ", np.sum(abs(res)))
return res, norbs
def vxc_pack(self, **kw):
"""
Computes the exchange-correlation matrix elements packed version (upper triangular)
Args:
sv : (System Variables), this must have arrays of coordinates and species, etc
Returns:
vxc,exc
"""
from pyscf.nao.m_xc_scalar_ni import xc_scalar_ni
from pyscf.nao.m_ao_matelem import ao_matelem_c
from pyscf.nao.m_pack2den import cp_block_pack_u, pack2den_u
#sv, dm, xc_code, deriv, kernel=None, ao_log=None, dtype=float64, **kvargs
sv = self
dm = kw['dm'] if 'dm' in kw else self.make_rdm1()
ao_log = kw['ao_log'] if 'ao_log' in kw else self.ao_log
xc_code = kw['xc_code'] if 'xc_code' in kw else self.xc_code
kw.pop('xc_code', None)
dtype = kw['dtype'] if 'dtype' in kw else self.dtype
aome = ao_matelem_c(ao_log.rr, ao_log.pp, sv, dm)
me = aome.init_one_set(ao_log)
atom2s = zeros((sv.natm + 1), dtype=int64)
for atom, sp in enumerate(sv.atom2sp):
atom2s[atom + 1] = atom2s[atom] + me.ao1.sp2norbs[sp]
sp2rcut = array([max(mu2rcut) for mu2rcut in me.ao1.sp_mu2rcut])
norbs = atom2s[-1]
nq, npk = (self.nspin - 1) * 2 + 1, norbs * (norbs + 1) // 2
kernel = kw['kernel'].reshape((nq, npk)) if 'kernel' in kw else zeros(
(nq, npk), dtype=dtype)
for atom1, [sp1, rv1, s1, f1] in enumerate(
zip(sv.atom2sp, sv.atom2coord, atom2s, atom2s[1:])):
for atom2, [sp2, rv2, s2, f2] in enumerate(
zip(sv.atom2sp, sv.atom2coord, atom2s, atom2s[1:])):
if atom2 > atom1: continue
if (sp2rcut[sp1] + sp2rcut[sp2])**2 <= sum((rv1 - rv2)**2):
continue
iab2block = xc_scalar_ni(me,
sp1,
rv1,
sp2,
rv2,
xc_code=xc_code,
**kw)
for i, ab2v in enumerate(iab2block):
cp_block_pack_u(ab2v, s1, f1, s2, f2, kernel[i], add=True)
return kernel
def comp_coulomb_pack(sv, ao_log=None, funct=coulomb_am, dtype=np.float64, **kvargs):
"""
Computes the matrix elements given by funct, for instance coulomb interaction
Args:
sv : (System Variables), this must have arrays of coordinates and species, etc
ao_log : description of functions (either orbitals or product basis functions)
Returns:
matrix elements for the whole system in packed form (lower triangular part)
"""
from pyscf.nao.m_ao_matelem import ao_matelem_c
from pyscf.nao.m_pack2den import ij2pack_l
aome = ao_matelem_c(sv.ao_log.rr, sv.ao_log.pp)
me = ao_matelem_c(sv.ao_log) if ao_log is None else aome.init_one_set(ao_log)
atom2s = np.zeros((sv.natm+1), dtype=np.int64)
for atom,sp in enumerate(sv.atom2sp): atom2s[atom+1]=atom2s[atom]+me.ao1.sp2norbs[sp]
norbs = atom2s[-1]
res = np.zeros(norbs*(norbs+1)//2, dtype=dtype)
for atom1,[sp1,rv1,s1,f1] in enumerate(zip(sv.atom2sp,sv.atom2coord,atom2s,atom2s[1:])):
#print("atom1 = {0}, rv1 = {1}".format(atom1, rv1))
for atom2,[sp2,rv2,s2,f2] in enumerate(zip(sv.atom2sp,sv.atom2coord,atom2s,atom2s[1:])):
if atom2>atom1: continue # skip
oo2f = funct(me,sp1,rv1,sp2,rv2, **kvargs)
if use_numba:
fill_triu_v2(oo2f, res, s1, f1, s2, f2, norbs)
else:
for i1 in range(s1,f1):
for i2 in range(s2, min(i1+1, f2)):
res[ij2pack_l(i1,i2,norbs)] = oo2f[i1-s1,i2-s2]
#print("number call = ", count)
#print("sum kernel: {0:.6f}".format(np.sum(abs(res))))
#np.savetxt("kernel_pyscf.txt", res)
#import sys
#sys.exit()
return res, norbs
def dipole_coo(sv, ao_log=None, funct=dipole_ni, **kvargs):
"""
Computes the dipole matrix and returns it in coo format (simplest sparse format to construct)
Args:
sv : (System Variables), this must have arrays of coordinates and species, etc
Returns:
overlap (real-space overlap) for the whole system
"""
from pyscf.nao.m_ao_matelem import ao_matelem_c
from scipy.sparse import coo_matrix
from numpy import array, int64, zeros
aome = ao_matelem_c(sv.ao_log.rr, sv.ao_log.pp)
me = aome.init_one_set(
sv.ao_log) if ao_log is None else aome.init_one_set(ao_log)
atom2s = zeros((sv.natm + 1), dtype=int64)
for atom, sp in enumerate(sv.atom2sp):
atom2s[atom + 1] = atom2s[atom] + me.ao1.sp2norbs[sp]
sp2rcut = array([max(mu2rcut) for mu2rcut in me.ao1.sp_mu2rcut])
nnz = 0
for sp1, rv1 in zip(sv.atom2sp, sv.atom2coord):
n1, rc1 = me.ao1.sp2norbs[sp1], sp2rcut[sp1]
for sp2, rv2 in zip(sv.atom2sp, sv.atom2coord):
if (rc1 + sp2rcut[sp2])**2 > ((rv1 - rv2)**2).sum():
nnz = nnz + n1 * me.ao1.sp2norbs[sp2]
irow, icol, data = zeros(nnz, dtype=int64), zeros(nnz, dtype=int64), zeros(
(3, nnz)) # Start to construct three coo matrices
inz = -1
for atom1, [sp1, rv1, s1, f1] in enumerate(
zip(sv.atom2sp, sv.atom2coord, atom2s, atom2s[1:])):
for atom2, [sp2, rv2, s2, f2] in enumerate(
zip(sv.atom2sp, sv.atom2coord, atom2s, atom2s[1:])):
if (sp2rcut[sp1] + sp2rcut[sp2])**2 <= sum((rv1 - rv2)**2):
continue
dd = funct(me, sp1, rv1, sp2, rv2, **kvargs)
for o1 in range(s1, f1):
for o2 in range(s2, f2):
inz = inz + 1
irow[inz], icol[inz], data[:, inz] = o1, o2, dd[:, o1 - s1,
o2 - s2]
norbs = atom2s[-1]
sh = (norbs, norbs)
rc = (irow, icol)
return coo_matrix((data[0], rc), shape=sh), coo_matrix(
(data[1], rc), shape=sh), coo_matrix((data[2], rc), shape=sh)
def vxc_lil(self, **kw):
"""
Computes the exchange-correlation matrix elements
Args:
sv : (System Variables), this must have arrays of coordinates and species, etc
Returns:
fxc,vxc,exc
"""
from pyscf.nao.m_xc_scalar_ni import xc_scalar_ni
from pyscf.nao.m_ao_matelem import ao_matelem_c
from scipy.sparse import lil_matrix
#dm, xc_code, deriv, ao_log=None, dtype=float64, **kvargs
dm = kw['dm'] if 'dm' in kw else self.make_rdm1()
kernel = kw['kernel'] if 'kernel' in kw else None
ao_log = kw['ao_log'] if 'ao_log' in kw else self.ao_log
xc_code = kw['xc_code'] if 'xc_code' in kw else self.xc_code
kw.pop('xc_code', None)
dtype = kw['dtype'] if 'dtype' in kw else self.dtype
aome = ao_matelem_c(self.ao_log.rr, self.ao_log.pp, self, dm)
me = aome.init_one_set(
self.ao_log) if ao_log is None else aome.init_one_set(ao_log)
atom2s = zeros((self.natm + 1), dtype=int64)
for atom, sp in enumerate(self.atom2sp):
atom2s[atom + 1] = atom2s[atom] + me.ao1.sp2norbs[sp]
lil = [
lil_matrix((atom2s[-1], atom2s[-1]), dtype=dtype)
for i in range((self.nspin - 1) * 2 + 1)
]
for atom1, [sp1, rv1, s1, f1] in enumerate(
zip(self.atom2sp, self.atom2coord, atom2s, atom2s[1:])):
for atom2, [sp2, rv2, s2, f2] in enumerate(
zip(self.atom2sp, self.atom2coord, atom2s, atom2s[1:])):
blk = xc_scalar_ni(me, sp1, rv1, sp2, rv2, xc_code=xc_code, **kw)
for i, b in enumerate(blk):
lil[i][s1:f1, s2:f2] = b[:, :]
return lil
def dipole_coo(sv, ao_log=None, funct=dipole_ni, **kvargs):
"""
Computes the dipole matrix and returns it in coo format (simplest sparse format to construct)
Args:
sv : (System Variables), this must have arrays of coordinates and species, etc
Returns:
overlap (real-space overlap) for the whole system
"""
from pyscf.nao.m_ao_matelem import ao_matelem_c
from scipy.sparse import coo_matrix
from numpy import array, int64, zeros
aome = ao_matelem_c(sv.ao_log.rr, sv.ao_log.pp)
me = aome.init_one_set(sv.ao_log) if ao_log is None else aome.init_one_set(ao_log)
atom2s = zeros((sv.natm+1), dtype=int64)
for atom,sp in enumerate(sv.atom2sp): atom2s[atom+1]=atom2s[atom]+me.ao1.sp2norbs[sp]
sp2rcut = array([max(mu2rcut) for mu2rcut in me.ao1.sp_mu2rcut])
nnz = 0
for sp1,rv1 in zip(sv.atom2sp,sv.atom2coord):
n1,rc1 = me.ao1.sp2norbs[sp1],sp2rcut[sp1]
for sp2,rv2 in zip(sv.atom2sp,sv.atom2coord):
if (rc1+sp2rcut[sp2])**2>((rv1-rv2)**2).sum() : nnz = nnz + n1*me.ao1.sp2norbs[sp2]
irow,icol,data = zeros(nnz, dtype=int64),zeros(nnz, dtype=int64),zeros((3,nnz)) # Start to construct three coo matrices
inz=-1
for atom1,[sp1,rv1,s1,f1] in enumerate(zip(sv.atom2sp,sv.atom2coord,atom2s,atom2s[1:])):
for atom2,[sp2,rv2,s2,f2] in enumerate(zip(sv.atom2sp,sv.atom2coord,atom2s,atom2s[1:])):
if (sp2rcut[sp1]+sp2rcut[sp2])**2<=sum((rv1-rv2)**2) : continue
dd = funct(me,sp1,rv1,sp2,rv2,**kvargs)
for o1 in range(s1,f1):
for o2 in range(s2,f2):
inz = inz+1
irow[inz],icol[inz],data[:,inz] = o1,o2,dd[:,o1-s1,o2-s2]
norbs = atom2s[-1]
sh = (norbs,norbs)
rc = (irow,icol)
return coo_matrix((data[0], rc), shape=sh),coo_matrix((data[1], rc), shape=sh),coo_matrix((data[2],rc), shape=sh)
l2S.fill(0.0)
for l3 in range(abs(l1 - l2), l1 + l2 + 1):
l2S[l3] = (f1f2_mom[:] * bessel_pp[l3, :]).sum(
) + f1f2_mom[0] * bessel_pp[l3, 0] / dkappa
cS.fill(0.0)
for m1 in range(-l1, l1 + 1):
for m2 in range(-l2, l2 + 1):
gc, m3 = self.get_gaunt(l1, -m1, l2, m2), m2 - m1
for l3ind, l3 in enumerate(
range(abs(l1 - l2), l1 + l2 + 1)):
if abs(m3) > l3: continue
cS[m1+_j,m2+_j] = cS[m1+_j,m2+_j] + l2S[l3]*ylm[ l3*(l3+1)+m3] *\
gc[l3ind] * (-1.0)**((3*l1+l2+l3)//2+m2)
self.c2r_(l1, l2, self.jmx, cS, rS, cmat)
oo2co[s1:f1, s2:f2] = rS[-l1 + _j:l1 + _j + 1,
-l2 + _j:l2 + _j + 1]
#sys.exit()
oo2co = oo2co * (4 * np.pi)**2 * self.interp_pp.dg_jt
return oo2co
if __name__ == '__main__':
from pyscf.nao.m_system_vars import system_vars_c
from pyscf.nao.m_ao_matelem import ao_matelem_c
sv = system_vars_c("siesta")
ra = np.array([0.0, 0.1, 0.2])
rb = np.array([0.0, 0.1, 0.0])
coulo = ao_matelem_c(sv.ao_log).coulomb_am(me, 0, ra, 0, rb)
def comp_vext_tem_pyth(self, ao_log=None, numba_parallel=True):
"""
Compute the external potential created by a moving charge
Python version
"""
def c2r_lm(conv, clm, clmm, m):
"""
clm: sph harmonic l and m
clmm: sph harmonic l and -m
convert from real to complex spherical harmonic
for an unique value of l and m
"""
rlm = 0.0
if m == 0:
rlm = conv._c2r[conv._j, conv._j] * clm
else:
rlm = conv._c2r[m+conv._j, m+conv._j]*clm +\
conv._c2r[m+conv._j, -m+conv._j]*clmm
if rlm.imag > 1e-10:
print(rlm)
raise ValueError("Non nul imaginary paert for c2r conversion")
return rlm.real
def get_index_lm(l, m):
"""
return the index of an array ordered as
[l=0 m=0, l=1 m=-1, l=1 m=0, l=1 m=1, ....]
"""
return (l + 1)**2 - 1 - l + m
warnings.warn("Obselete routine use comp_vext_tem")
if use_numba:
get_time_potential = nb.jit(
nopython=True, parallel=numba_parallel)(get_tem_potential_numba)
V_time = np.zeros((self.time.size), dtype=np.complex64)
aome = ao_matelem_c(self.ao_log.rr, self.ao_log.pp)
me = ao_matelem_c(
self.ao_log) if ao_log is None else aome.init_one_set(ao_log)
atom2s = np.zeros((self.natm + 1), dtype=np.int64)
for atom, sp in enumerate(self.atom2sp):
atom2s[atom + 1] = atom2s[atom] + me.ao1.sp2norbs[sp]
R0 = self.vnorm * self.time[0] * self.vdir + self.beam_offset
rr = self.ao_log.rr
dr = (np.log(rr[-1]) - np.log(rr[0])) / (rr.size - 1)
dt = self.time[1] - self.time[0]
dw = self.freq_symm[1] - self.freq_symm[0]
wmin = self.freq_symm[0]
tmin = self.time[0]
nff = self.freq.size
ub = self.freq_symm.size // 2 - 1
l2m = [] # list storing m value to corresponding l
fact_fft = np.exp(-1.0j * self.freq_symm[ub:ub + nff] * tmin)
pre_fact = dt * np.exp(-1.0j * wmin * (self.time - tmin))
for l in range(me.jmx + 1):
lm = []
for m in range(-l, l + 1):
lm.append(m)
l2m.append(np.array(lm))
for atm, sp in enumerate(self.atom2sp):
rcut = self.ao_log.sp2rcut[sp]
center = self.atom2coord[atm, :]
rmax = find_nearrest_index(rr, rcut)
si = atom2s[atm]
fi = atom2s[atm + 1]
for mu, l in enumerate(self.pb.prod_log.sp_mu2j[sp]):
s = self.pb.prod_log.sp_mu2s[sp][mu]
f = self.pb.prod_log.sp_mu2s[sp][mu + 1]
fr_val = self.pb.prod_log.psi_log[sp][mu, :]
inte1 = np.sum(fr_val[0:rmax + 1] * rr[0:rmax + 1]**(l + 2) *
rr[0:rmax + 1] * dr)
for k in range(s, f):
V_time.fill(0.0)
m = l2m[l][k - s]
ind_lm = get_index_lm(l, m)
ind_lmm = get_index_lm(l, -m)
if use_numba:
get_time_potential(self.time, R0, self.vnorm, self.vdir,
center, rcut, inte1, rr, dr, fr_val,
me._c2r, l, m, me._j, ind_lm, ind_lmm,
V_time)
else:
for it, t in enumerate(self.time):
R_sub = R0 + self.vnorm * self.vdir * (
t - self.time[0]) - center
norm = np.sqrt(np.dot(R_sub, R_sub))
if norm > rcut:
I1 = inte1 / (norm**(l + 1))
I2 = 0.0
else:
rsub_max = find_nearrest_index(rr, norm)
I1 = np.sum(fr_val[0:rsub_max + 1] *
rr[0:rsub_max + 1]**(l + 2) *
rr[0:rsub_max + 1])
I2 = np.sum(fr_val[rsub_max + 1:] *
rr[rsub_max + 1:] /
(rr[rsub_max + 1:]**(l - 1)))
I1 = I1 * dr / (norm**(l + 1))
I2 = I2 * (norm**l) * dr
clm_tem = csphar(R_sub, l)
clm = (4 * np.pi /
(2 * l + 1)) * clm_tem[ind_lm] * (I1 + I2)
clmm = (4 * np.pi /
(2 * l + 1)) * clm_tem[ind_lmm] * (I1 + I2)
rlm = c2r_lm(me, clm, clmm, m)
V_time[it] = rlm + 0.0j
V_time *= pre_fact
FT = fft(V_time)
self.V_freq[:, si + k] = FT[ub:ub + nff] * fact_fft
for mu1,l1,s1,f1 in self.ao1.sp2info[sp1]:
f1f2_mom = self.ao1.psi_log_mom[sp2][mu2,:] * self.ao1.psi_log_mom[sp1][mu1,:]
l2S.fill(0.0)
for l3 in range( abs(l1-l2), l1+l2+1):
l2S[l3] = (f1f2_mom[:]*bessel_pp[l3,:]).sum() + f1f2_mom[0]*bessel_pp[l3,0]/dkappa
cS.fill(0.0)
for m1 in range(-l1,l1+1):
for m2 in range(-l2,l2+1):
gc,m3 = self.get_gaunt(l1,-m1,l2,m2), m2-m1
for l3ind,l3 in enumerate(range(abs(l1-l2),l1+l2+1)):
if abs(m3) > l3 : continue
cS[m1+_j,m2+_j] = cS[m1+_j,m2+_j] + l2S[l3]*ylm[ l3*(l3+1)+m3] *\
gc[l3ind] * (-1.0)**((3*l1+l2+l3)//2+m2)
self.c2r_( l1,l2, self.jmx,cS,rS,cmat)
oo2co[s1:f1,s2:f2] = rS[-l1+_j:l1+_j+1,-l2+_j:l2+_j+1]
#sys.exit()
oo2co = oo2co * (4*np.pi)**2 * self.interp_pp.dg_jt
return oo2co
if __name__ == '__main__':
from pyscf.nao.m_system_vars import system_vars_c
from pyscf.nao.m_prod_log import prod_log_c
from pyscf.nao.m_ao_matelem import ao_matelem_c
sv = system_vars_c("siesta")
ra = np.array([0.0, 0.1, 0.2])
rb = np.array([0.0, 0.1, 0.0])
coulo = ao_matelem_c(sv.ao_log).coulomb_am(me, 0, ra, 0, rb)
def overlap_coo(sv, ao_log=None, funct=overlap_ni, ao_log2=None, **kw):
"""
Computes the overlap matrix and returns it in coo format (simplest sparse format to construct)
Args:
sv : (System Variables), this must have arrays of coordinates and species, etc
Returns:
overlap (real-space overlap) for the whole system
"""
from pyscf.nao.m_ao_matelem import ao_matelem_c
from scipy.sparse import coo_matrix
from numpy import array, float64, int64, zeros
aome = ao_matelem_c(sv.ao_log.rr, sv.ao_log.pp)
if ao_log is None and ao_log2 is None:
me = aome.init_one_set(sv.ao_log)
elif ao_log is not None and ao_log2 is None:
me = aome.init_one_set(ao_log)
elif ao_log is None and ao_log2 is not None:
me = aome.init_one_set(ao_log2)
else:
me = aome.init_two_sets(ao_log, ao_log2)
a2s1 = zeros((sv.natm + 1), dtype=int)
a2s2 = zeros((sv.natm + 1), dtype=int)
for atom, sp in enumerate(sv.atom2sp):
a2s1[atom + 1] = a2s1[atom] + me.ao1.sp2norbs[sp]
a2s2[atom + 1] = a2s2[atom] + me.ao2.sp2norbs[sp]
sp2rcut1 = array([max(mu2rcut) for mu2rcut in me.ao1.sp_mu2rcut])
sp2rcut2 = array([max(mu2rcut) for mu2rcut in me.ao2.sp_mu2rcut])
nnz = 0
for sp1, rv1 in zip(sv.atom2sp, sv.atom2coord):
n1, rc1 = me.ao1.sp2norbs[sp1], sp2rcut1[sp1]
for sp2, rv2 in zip(sv.atom2sp, sv.atom2coord):
n2, rc2 = me.ao2.sp2norbs[sp2], sp2rcut2[sp2]
if (rc1 + rc2)**2 > ((rv1 - rv2)**2).sum(): nnz = nnz + n1 * n2
irow, icol, data = zeros(nnz, dtype=int64), zeros(nnz, dtype=int64), zeros(
nnz) # Start to construct coo matrix
inz = -1
for atom1, [sp1, rv1, s1,
f1] in enumerate(zip(sv.atom2sp, sv.atom2coord, a2s1,
a2s1[1:])):
for atom2, [sp2, rv2, s2, f2] in enumerate(
zip(sv.atom2sp, sv.atom2coord, a2s2, a2s2[1:])):
if (sp2rcut1[sp1] + sp2rcut2[sp2])**2 <= sum((rv1 - rv2)**2):
continue
oo = funct(me, sp1, rv1, sp2, rv2, **kw)
for o1 in range(s1, f1):
for o2 in range(s2, f2):
inz = inz + 1
irow[inz], icol[inz], data[inz] = o1, o2, oo[o1 - s1,
o2 - s2]
norbs1 = a2s1[-1]
norbs2 = a2s2[-1]
return coo_matrix((data, (irow, icol)), shape=(norbs1, norbs2))
def vxc_pack(sv,
dm,
xc_code,
deriv,
kernel=None,
ao_log=None,
dtype=float64,
**kvargs):
"""
Computes the exchange-correlation matrix elements packed version (upper triangular)
Args:
sv : (System Variables), this must have arrays of coordinates and species, etc
Returns:
vxc,exc
"""
from pyscf.nao.m_xc_scalar_ni import xc_scalar_ni
from pyscf.nao.m_ao_matelem import ao_matelem_c
from pyscf.nao.m_pack2den import triu_indices
aome = ao_matelem_c(sv.ao_log.rr, sv.ao_log.pp, sv, dm)
me = aome.init_one_set(
sv.ao_log) if ao_log is None else aome.init_one_set(ao_log)
atom2s = zeros((sv.natm + 1), dtype=int64)
for atom, sp in enumerate(sv.atom2sp):
atom2s[atom + 1] = atom2s[atom] + me.ao1.sp2norbs[sp]
sp2rcut = array([max(mu2rcut) for mu2rcut in me.ao1.sp_mu2rcut])
norbs = atom2s[-1]
ind = triu_indices(norbs)
if kernel is None:
kernel = zeros(norbs * (norbs + 1) // 2, dtype=dtype)
for atom1, [sp1, rv1, s1, f1] in enumerate(
zip(sv.atom2sp, sv.atom2coord, atom2s, atom2s[1:])):
for atom2, [sp2, rv2, s2, f2] in enumerate(
zip(sv.atom2sp, sv.atom2coord, atom2s, atom2s[1:])):
if (sp2rcut[sp1] + sp2rcut[sp2])**2 <= sum((rv1 - rv2)**2):
continue
xc = xc_scalar_ni(me, sp1, rv1, sp2, rv2, xc_code, deriv,
**kvargs)
if use_numba:
fill_triu(xc, ind, kernel, s1, f1, s2, f2)
else:
for i1 in range(s1, f1):
for i2 in range(s2, f2):
if ind[i1, i2] >= 0:
kernel[ind[i1, i2]] = xc[i1 - s1, i2 - s2]
return kernel
else:
if kernel.size != int(norbs * (norbs + 1) // 2):
raise ValueError("wrong dimension for kernel")
for atom1, [sp1, rv1, s1, f1] in enumerate(
zip(sv.atom2sp, sv.atom2coord, atom2s, atom2s[1:])):
for atom2, [sp2, rv2, s2, f2] in enumerate(
zip(sv.atom2sp, sv.atom2coord, atom2s, atom2s[1:])):
if (sp2rcut[sp1] + sp2rcut[sp2])**2 <= sum((rv1 - rv2)**2):
continue
xc = xc_scalar_ni(me, sp1, rv1, sp2, rv2, xc_code, deriv,
**kvargs)
if use_numba:
fill_triu(xc, ind, kernel, s1, f1, s2, f2, add=True)
else:
for i1 in range(s1, f1):
for i2 in range(s2, f2):
if ind[i1, i2] >= 0:
kernel[ind[i1, i2]] += xc[i1 - s1, i2 - s2]
