def ee_energy(configs):
ne = configs.shape[1]
if ne == 1:
return np.zeros(configs.shape[0])
ee = np.zeros(configs.shape[0])
d = RawDistance()
ee, ij = d.dist_matrix(configs)
ee = np.linalg.norm(ee, axis=2)
return np.sum(1.0 / ee, axis=1)
def ii_energy(mol):
ei = 0.0
d = RawDistance()
rij, ij = d.dist_matrix(mol.atom_coords()[np.newaxis, :, :])
if len(ij) == 0:
return np.array([0.0])
rij = np.linalg.norm(rij, axis=2)[0, :]
iitot = 0
c = mol.atom_charges()
for (i, j), r in zip(ij, rij):
iitot += c[i] * c[j] / r
return iitot
def __init__(self,mol,a_basis=None,b_basis=None,dist=RawDistance() ):
"""
Args:
mol : a pyscf molecule object
a_basis : list of func3d objects that comprise the electron-ion basis
b_basis : list of func3d objects that comprise the electron-electron basis
dist: a distance calculator
"""
if b_basis is None:
nexpand=5
self.b_basis=[GaussianFunction(0.2*2**n) for n in range(1,nexpand+1)]
else:
nexpand=len(b_basis)
self.b_basis=b_basis
if a_basis is None:
aexpand=4
self.a_basis=[GaussianFunction(0.2*2**n) for n in range(1,aexpand+1)]
else:
aexpand=len(a_basis)
self.a_basis=a_basis
self.parameters={}
self._nelec=np.sum(mol.nelec)
self._mol=mol
self._dist=dist
self.parameters['bcoeff']=np.zeros((nexpand, 3))
self.parameters['acoeff']=np.zeros((aexpand, 2))
def ecp_ea(mol, configs, wf, e, atom, threshold):
"""
Returns the ECP value between electron e and atom at, local+nonlocal.
"""
from pyqmc.distance import RawDistance
nconf = configs.shape[0]
ecp_val = np.zeros(nconf, dtype=complex if wf["iscomplex"] else float)
at_name, apos = atom
apos = np.asarray(apos)
r_ea_vec = RawDistance().dist_i(apos, configs[:, e, :]).reshape((-1, 3))
r_ea = np.linalg.norm(r_ea_vec, axis=-1)
l_list, v_l = get_v_l(mol, at_name, r_ea)
mask, prob = ecp_mask(v_l, threshold)
masked_v_l = v_l[mask]
masked_v_l[:, :-1] /= prob[mask, np.newaxis]
# Use masked objects internally
r_ea = r_ea[mask]
r_ea_vec = r_ea_vec[mask]
P_l, r_ea_i = get_P_l(r_ea, r_ea_vec, l_list)
# Note: epos_rot is not just apos+r_ea_i because of the boundary;
# positions of the samples are relative to the electron, not atom.
epos_rot = (configs[mask, e, :] - r_ea_vec)[:, np.newaxis] + r_ea_i
# Expand externally
expanded_epos_rot = np.zeros((nconf, P_l.shape[1], 3))
expanded_epos_rot[mask] = epos_rot
epos = expanded_epos_rot
ratio = wf["testvalue"](configs, e, epos, mask)
# Compute local and non-local parts
ecp_val[mask] = np.einsum("ij,ik,ijk->i", ratio, masked_v_l, P_l)
ecp_val += v_l[:, -1] # local part
return ecp_val
def __init__(self, configs):
self.configs = configs
self.dist = RawDistance()