Skip to content

Commit

Permalink
add more debugs on SO3
Browse files Browse the repository at this point in the history
  • Loading branch information
@qzhu2017
qzhu2017 committed Oct 3, 2024
1 parent 821ef26 commit 0d2a100
Show file tree
Hide file tree
Showing 2 changed files with 75 additions and 57 deletions.
113 changes: 57 additions & 56 deletions pyxtal/lego/SO3.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -191,48 +191,49 @@ def compute_dpdr(self, atoms, atom_ids=None):
p_list = np.zeros((self.natoms, self.ncoefs), dtype=np.float64)
dp_list = np.zeros((self.natoms, self.natoms, self.ncoefs, 3), dtype=np.float64)

# get expansion coefficients and derivatives
cs, dcs = compute_dcs(self.neighborlist, self.nmax, self.lmax, self.rcut, self.alpha, self._cutoff_function)
if len(self.neighborlist) > 0:
# get expansion coefficients and derivatives
cs, dcs = compute_dcs(self.neighborlist, self.nmax, self.lmax, self.rcut, self.alpha, self._cutoff_function)

# weight cs and dcs
cs *= self.atomic_weights[:, np.newaxis, np.newaxis, np.newaxis]
dcs *= self.atomic_weights[:, np.newaxis, np.newaxis, np.newaxis, np.newaxis]
cs = np.einsum('inlm,l->inlm', cs, self.norm)
dcs = np.einsum('inlmj,l->inlmj', dcs, self.norm)
#print('cs, dcs', self.neighbor_indices, cs.shape, dcs.shape)
# weight cs and dcs
cs *= self.atomic_weights[:, np.newaxis, np.newaxis, np.newaxis]
dcs *= self.atomic_weights[:, np.newaxis, np.newaxis, np.newaxis, np.newaxis]
cs = np.einsum('inlm,l->inlm', cs, self.norm)
dcs = np.einsum('inlmj,l->inlmj', dcs, self.norm)
#print('cs, dcs', self.neighbor_indices, cs.shape, dcs.shape)

# Assign cs and dcs to P and dP
# cs: (N_ij, n, l, m) => P (N_i, N_des)
# dcs: (N_ij, n, l, m, 3) => dP (N_i, N_j, N_des, 3)
# (n, l, m) needs to be merged to 1 dimension
# Assign cs and dcs to P and dP
# cs: (N_ij, n, l, m) => P (N_i, N_des)
# dcs: (N_ij, n, l, m, 3) => dP (N_i, N_j, N_des, 3)
# (n, l, m) needs to be merged to 1 dimension

for i in range(len(atoms)):
# find atoms for which i is the center
centers = self.neighbor_indices[:, 0] == i
for i in range(len(atoms)):
# find atoms for which i is the center
centers = self.neighbor_indices[:, 0] == i

if len(self.neighbor_indices[centers]) > 0:
# total up the c array for the center atom
ctot = cs[centers].sum(axis=0) #(n, l, m)
if len(self.neighbor_indices[centers]) > 0:
# total up the c array for the center atom
ctot = cs[centers].sum(axis=0) #(n, l, m)

# power spectrum P = c*c_conj
# eq_3 (n, n', l) eliminate m
P = np.einsum('ijk, ljk->ilj', ctot, np.conj(ctot)).real
p_list[i] = P[self.tril_indices].flatten()
# power spectrum P = c*c_conj
# eq_3 (n, n', l) eliminate m
P = np.einsum('ijk, ljk->ilj', ctot, np.conj(ctot)).real
p_list[i] += P[self.tril_indices].flatten()

# gradient of P for each neighbor, eq_26
# (N_ijs, n, n', l, 3)
# dc * c_conj + c * dc_conj
dP = np.einsum('wijkn,ljk->wiljn', dcs[centers], np.conj(ctot))
dP += np.conj(np.transpose(dP, axes=[0, 2, 1, 3, 4]))
dP = dP.real
# gradient of P for each neighbor, eq_26
# (N_ijs, n, n', l, 3)
# dc * c_conj + c * dc_conj
dP = np.einsum('wijkn,ljk->wiljn', dcs[centers], np.conj(ctot))
dP += np.conj(np.transpose(dP, axes=[0, 2, 1, 3, 4]))
dP = dP.real

#print("shape of P/dP", P.shape, dP.shape)#; import sys; sys.exit()
#print("shape of P/dP", P.shape, dP.shape)#; import sys; sys.exit()

# QZ: to check
ijs = self.neighbor_indices[centers]
for _id, j in enumerate(ijs[:, 1]):
dp_list[i, j, :, :] += dP[_id][self.tril_indices].flatten().reshape(self.ncoefs, 3)
dp_list[i, i, :, :] -= dP[_id][self.tril_indices].flatten().reshape(self.ncoefs, 3)
# QZ: to check
ijs = self.neighbor_indices[centers]
for _id, j in enumerate(ijs[:, 1]):
dp_list[i, j, :, :] += dP[_id][self.tril_indices].flatten().reshape(self.ncoefs, 3)
dp_list[i, i, :, :] -= dP[_id][self.tril_indices].flatten().reshape(self.ncoefs, 3)

return dp_list, p_list

Expand All @@ -242,7 +243,6 @@ def compute_dpdr_5d(self, atoms):
Args:
atoms: ase atoms object
atom_ids: optional list of atomic indices
Returns:
dpdr array (N, N, M, 3, 27) and p array (N, M)
Expand Down Expand Up @@ -270,31 +270,31 @@ def compute_dpdr_5d(self, atoms):
for i in range(len(atoms)):
# find atoms for which i is the center
pair_ids = neigh_ids[self.neighbor_indices[:, 0] == i]

# loop over each pair
for pair_id in pair_ids:
(_, j, x, y, z) = self.neighbor_indices[pair_id]
# map from (x, y, z) to (0, 27)
cell_id = (x+1) * 9 + (y+1) * 3 + z + 1

if len(pair_ids) > 0:
ctot = cs[pair_ids].sum(axis=0) #(n, l, m)
# power spectrum P = c*c_conj
# eq_3 (n, n', l) eliminate m
P = np.einsum('ijk, ljk->ilj', cs[pair_id], np.conj(cs[pair_id])).real
p_list[i] = P[self.tril_indices].flatten()
P = np.einsum('ijk, ljk->ilj', ctot, np.conj(ctot)).real
p_list[i] += P[self.tril_indices].flatten()

# gradient of P for each neighbor, eq_26
# (N_ijs, n, n', l, 3)
# dc * c_conj + c * dc_conj
dP = np.einsum('ijkn, ljk->iljn', dcs[pair_id], np.conj(cs[pair_id]))
dP += np.conj(np.transpose(dP, axes=[1, 0, 2, 3]))
dP = dP.real[self.tril_indices].flatten().reshape(self.ncoefs, 3)
#print(cs[pair_id].shape, dcs[pair_id].shape, dP.shape)
# loop over each pair
for pair_id in pair_ids:
(_, j, x, y, z) = self.neighbor_indices[pair_id]
# map from (x, y, z) to (0, 27)
cell_id = (x+1) * 9 + (y+1) * 3 + z + 1

dp_list[i, j, :, :, cell_id] += dP
dp_list[i, i, :, :, cell_id] -= dP
# gradient of P for each neighbor, eq_26
# (N_ijs, n, n', l, 3)
# dc * c_conj + c * dc_conj
dP = np.einsum('ijkn, ljk->iljn', dcs[pair_id], np.conj(ctot))
dP += np.conj(np.transpose(dP, axes=[1, 0, 2, 3]))
dP = dP.real[self.tril_indices].flatten().reshape(self.ncoefs, 3)
#print(cs[pair_id].shape, dcs[pair_id].shape, dP.shape)

return dp_list, p_list
dp_list[i, j, :, :, cell_id] += dP
dp_list[i, i, :, :, cell_id] -= dP

return dp_list, p_list


def calculate(self, atoms, atom_ids=None, derivative=False):
Expand Down Expand Up @@ -696,5 +696,6 @@ def compute_dcs(pos, nmax, lmax, rcut, alpha, cutoff):
print(f)
print('x', x['x'])
#print('dxdr', x['dxdr'])
print('calculation time {}'.format(start2-start1))
print(f.compute_p(test))
p = f.compute_p(test); print('from P', p)
dp, p = f.compute_dpdr(test); print('from dP', p)
dp, p = f.compute_dpdr_5d(test); print('from dP5d', p)
19 changes: 18 additions & 1 deletion tests/test_SO3.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -45,7 +45,7 @@ def get_dPdR_xtal(xtal, nmax, lmax, rc, eps):

# Descriptors Parameters
eps = 1e-8
rc = 4.00
rc = 2.80
nmax, lmax = 2, 2

# NaCl cluster
Expand Down Expand Up @@ -89,6 +89,23 @@ def test_dPdR_graphite(self):
c.from_prototype('graphite')
get_dPdR_xtal(c.to_ase(), nmax, lmax, rc, eps)

def test_dPdR_random(self):
x = [ 7.952, 2.606, 0.592, 0.926, 0.608, 0.307]
c = pyxtal()
c.from_spg_wps_rep(179, ['6a', '6a', '6a', '6a'], x)
get_dPdR_xtal(c.to_ase(), nmax, lmax, rc, eps)

def test_dPdR_random_P(self):
x = [ 7.952, 2.606, 0.592, 0.926, 0.608, 0.307]
c = pyxtal()
c.from_spg_wps_rep(179, ['6a', '6a', '6a', '6a'], x)
atoms = c.to_ase()
f = SO3(nmax=nmax, lmax=lmax, rcut=rc)
p0 = f.compute_p(atoms)
_, p1 = f.compute_dpdr(atoms)
_, p2 = f.compute_dpdr_5d(atoms)
assert(np.allclose(p0, p1, atol=1e-3))
assert(np.allclose(p0, p2, atol=1e-3))

if __name__ == "__main__":
unittest.main()

0 comments on commit 0d2a100

Please sign in to comment.