Source code for pyxtal_ff.descriptors.EAD

import math
import numpy as np
from ase.neighborlist import NeighborList
from .cutoff import Cutoff


[docs]class EAD: """EAD is an atom-centered descriptor that is inspired by Embedded Atom method (EAM). The EAM utilizes the orbital-dependent density components. The orbital-dependent component consists of a set of local atomic density descriptions. The functional form of EAD is consistent with: Zhang, Y., et. al. (2019). The Journal of Physical Chemistry Letters, 10(17), 4962-4967. Parameters ---------- parameters: dict The user-defined parameters for component of local atomic density descriptions. i.e. {'L': 2, 'eta': [0.36], 'Rs': [1.0]} Rc: float The EAD will be calculated within this radius. derivative: bool If True, calculate the derivative of EAD. stress: bool If True, calculate the virial stress contribution of EAD. """ def __init__(self, parameters, Rc=5., derivative=True, stress=False, cutoff='cosine'): self._type = 'EAD' self.dsize = int(1) self.parameters = {} # Set up keywords keywords = ['L', 'eta', 'Rs'] for k, v in parameters.items(): if k not in keywords: msg = f"{k} is not a valid key. "\ f"Choose from {keywords}" raise NotImplementedError(msg) else: if k == 'L': self.dsize *= (v+1) self.parameters[k] = v else: self.dsize *= len(v) self.parameters[k] = np.array(v) self.parameters['cutoff'] = cutoff self.Rc = Rc self.derivative = derivative self.stress = stress def __str__(self): s = "EAMD descriptor with Cutoff: {:6.3f}\n".format(self.Rc) for key in self.parameters.keys(): s += " {:s}: ".format(key) vals = self.parameters[key] if key in ['eta']: for val in vals: s += "{:8.3f}, ".format(val) elif key in ['Rs']: for val in vals: s += "{:6.3f}".format(val) else: s += "{:2d}".format(vals) s += "\n" return s def __repr__(self): return str(self)
[docs] def load_from_dict(self, dict0): self.parameters = dict0["parameters"] self.Rc = dict0["Rc"] self.derivative = dict0["derivative"] self.stress = dict0["stress"]
[docs] def save_dict(self): """ save the model as a dictionary in json """ dict = { "parameters": self.parameters, "rcut": self.rcut, "derivative": self.derivative, "stress": self.stress, } return dict
[docs] def calculate(self, crystal, ids=None): """Calculate and return the EAD. Parameters ---------- crystal: object ASE Structure object. ids: list A list of the centered atoms to be computed if None, all atoms will be considered Returns ------- d: dict The user-defined EAD that represent the crystal. d = {'x': [N, d], 'dxdr': [N, m, d, 3], 'rdxdr': [N, m, d, 3, 3], 'elements': list of elements} """ self.crystal = crystal self.total_atoms = len(crystal) # total atoms in the unit cell vol = crystal.get_volume() rc = [self.Rc/2.]*self.total_atoms neighbors = NeighborList(rc, self_interaction=False, bothways=True, skin=0.) neighbors.update(crystal) unique_N = 0 for i in range(self.total_atoms): indices, offsets = neighbors.get_neighbors(i) ith = 0 if i not in indices: ith += 1 unique_N += len(np.unique(indices))+ith # +1 is for i # Make numpy array here self.d = {'x': np.zeros([self.total_atoms, self.dsize]), 'elements': []} if self.derivative: self.d['dxdr'] = np.zeros([unique_N, self.dsize, 3]) self.d['seq'] = np.zeros([unique_N, 2], dtype=int) if self.stress: self.d['rdxdr'] = np.zeros([unique_N, self.dsize, 3, 3]) seq_count = 0 if ids is None: ids = range(len(crystal)) for i in ids: #range(self.total_atoms): element = crystal.get_chemical_symbols()[i] indices, offsets = neighbors.get_neighbors(i) Z = [] # atomic numbers of neighbors assert len(indices) > 0, \ f"There's no neighbor for this structure at Rc = {self.Rc} A." Ri = crystal.get_positions()[i] total_neighbors = len(indices) Rj = np.zeros([total_neighbors, 3]) IDs = np.zeros(total_neighbors, dtype=int) count = 0 for j, offset in zip(indices, offsets): Rj[count, :] = crystal.positions[j] + np.dot(offset, crystal.get_cell()) IDs[count] = j Z.append(crystal[j].number) count += 1 Z = np.array(Z) Rij = Rj - Ri Dij = np.sqrt(np.sum(Rij**2, axis=1)) d = calculate_eamd(i, self.total_atoms, Ri, Rij, Dij, Z, IDs, self.Rc, self.parameters, self.derivative, self.stress) self.d['x'][i] = d['x'] if self.derivative: n_seq = len(d['seq']) self.d['dxdr'][seq_count:seq_count+n_seq] = d['dxdr'] self.d['seq'][seq_count:seq_count+n_seq] = d['seq'] if self.stress: self.d['rdxdr'][seq_count:seq_count+n_seq] = d['rdxdr']/vol if self.derivative: seq_count += n_seq self.d['elements'].append(element) return self.d
[docs]def calculate_eamd(i, m, ri, rij, dij, Z, IDs, Rc, parameters, derivative, stress): """ Calculate the EAD for a center atom i. Parameters ---------- i: int The i-th atom center. m: int The total atoms in the crystal unit cell. ri: array [3] The position of atom i. rij: array [j, 3] The vector distances of atom i to neighbors j. dij: array [j] The array of distances of i-th center atom. Z: array [j] The atomic numbers of neighbors. IDs: int array [j] The indices of neighbors centering about atom i. Rc: float The cutoff radius. parameters: dict Rs: float array (d1) The shift from the center of the Gaussian-type orbitals. etas: float array (d2) The width of the Gaussian-type orbitals. L: int (d3) The total orbital angular momentum. cutoff: str The cutoff function. derivative: If True, calculate the derivative of EAD. stress: bool If True, calculate the virial stress contribution of EAD. Returns ------- Dict of EAD descriptors with its derivative and stress contribution. """ l_index = [1, 4, 10, 20] normalize = 1 / np.sqrt(np.array([1, 1, 1, 1, 1, 1, 1, 2, 2, 2, 1, 2, 2, 2, 2, 2, 2, 6, 6, 6])) Rs = parameters['Rs'] # d1 etas = parameters['eta'] # d2 L = parameters['L'] # d3 cutoff = Cutoff(parameters['cutoff']) d1, d2, d3, j = len(Rs), len(etas), L+1, len(dij) ij_list = i * np.ones([len(IDs), 2], dtype=int) ij_list[:, 1] = IDs unique_js = np.unique(IDs) if i not in unique_js: unique_js = np.append(i, unique_js) unique_js.sort() seq = i*np.ones([len(unique_js), 2], dtype=int) seq[:, 1] = unique_js uN = len(unique_js) _i = np.where(unique_js==i)[0][0] term1, d_term1, i_d_term1, j_d_term1 = get_xyz(unique_js, rij, ij_list, L, derivative=derivative) # [j, D3], [j, uN, 3, l] d0 = dij - Rs[:, np.newaxis] d02 = d0 ** 2 # [d1, j] fc = cutoff.calculate(dij, Rc) # [j] cj_cutoff = Z * fc # [j] term2_1 = np.exp(np.einsum('i,jk->ijk', -etas, d02)) # [d2, d1, j] term2 = np.einsum('ijk,k->ijk', term2_1, cj_cutoff) # [d2, d1, j] term = np.einsum('ij,kli->jkl', term1, term2) # [D3, d2, d1] if derivative: dterm0 = np.einsum('k, ijk->ijk', Z, term2_1).reshape([d1*d2, j]) # [d2*d1, j] dterm11 = np.einsum('ij, j->ij', dterm0, fc).reshape([d1*d2, j]) # [d2*d1, j] dterm1 = np.einsum('ij, jklm->jmilk', dterm11, d_term1) # [j, D3, d2*d1, 3, uN] i_dterm1 = np.einsum('ij, jklm->jmilk', dterm11, i_d_term1) j_dterm1 = np.einsum('ij, jklm->jmilk', dterm11, j_d_term1) dterm20 = np.einsum('ij, ki->jki', term1, dterm0) # [D3, d2*d1, j] dterm21 = cutoff.calculate_derivative(dij, Rc) # [j] _dterm22 = np.einsum('ij,j->ij', d0, fc) # [d1, j] dterm22 = 2 * np.einsum('i,jk->ijk', etas, _dterm22) # [d2, d1, j] dterm23 = (dterm21 - dterm22).reshape([d2*d1, j]) # [d2*d1, j] dterm24 = np.einsum('ijk, jk->ijk', dterm20, dterm23) # [D3, d2*d1, j] dRij_dRm = np.zeros([j, 3, uN]) i_dRij_dRm = np.zeros([j, 3, uN]) j_dRij_dRm = np.zeros([j, 3, uN]) for mm, _m in enumerate(unique_js): mm_list = _m * np.ones([j, 1], dtype=int) dRij_dRm[:,:,mm], i_dRij_dRm[:,:,mm], j_dRij_dRm[:,:,mm] = \ dRij_dRm_norm(rij, np.hstack((ij_list, mm_list))) # [j, 3, uN] dterm2 = np.einsum('ijk, klm->kijlm', dterm24, dRij_dRm) # [j, D3, d2*d1, 3, uN] i_dterm2 = np.einsum('ijk, klm->kijlm', dterm24, i_dRij_dRm) # [j, D3, d2*d1, 3, uN] j_dterm2 = np.einsum('ijk, klm->kijlm', dterm24, j_dRij_dRm) # [j, D3, d2*d1, 3, uN] dphi_dRm = dterm1 + dterm2 # [j, D3, d2*d1, 3, uN] i_dphi_dRm = i_dterm1 + i_dterm2 j_dphi_dRm = j_dterm1 + j_dterm2 dterm = np.einsum('ij, hijkl->ijkl', term.reshape([term.shape[0], d2*d1]), dphi_dRm) # [D3, d2*d1, 3, uN] if stress: _RDXDR = np.zeros([term.shape[0], d2*d1, 3, uN, 3]) # [D3, d2*d1, 3, uN, 3] for count, ij in enumerate(ij_list): _j = np.where(unique_js==ij[1])[0][0] i_tmp = i_dphi_dRm[count, :, :, :] j_tmp = j_dphi_dRm[count, :, :, :] _RDXDR[:, :, :, _i, :] += np.einsum('ijk,l->ijkl', i_tmp[:,:,:,_i], ri) _RDXDR[:, :, :, _j, :] += np.einsum('ijk,l->ijkl', j_tmp[:,:,:,_j], rij[count]+ri) sterm = np.einsum('ij, ijklm->ijklm', term.reshape([term.shape[0], d2*d1]), _RDXDR) count = 0 x = np.zeros([d3*d2*d1]) # [d3*d2*d1] dxdr, rdxdr = None, None if derivative: dxdr = np.zeros([uN, d1*d2*d3, 3]) # [uN, d3*d2*d1, 3] if stress: rdxdr = np.zeros([uN, d1*d2*d3, 3, 3]) for l in range(L+1): Rc2l = Rc**(2*l) L_fac = math.factorial(l) x[count:count+d1*d2] = L_fac * np.einsum('ijk->jk', term[:l_index[l]] ** 2).ravel() / Rc2l if derivative: dxdr[:, count:count+d1*d2, :] = 2 * L_fac * np.einsum('ijkl->ljk', dterm[:l_index[l]]) / Rc2l if stress: rdxdr[:, count:count+d1*d2, :, :] = 2 * L_fac * np.einsum('ijklm->ljkm', sterm[:l_index[l]]) / Rc2l count += d1*d2 return {'x': x, 'dxdr': dxdr, 'rdxdr': rdxdr, 'seq': seq}
[docs]def get_xyz(unique_js, rij, ij_list, L, derivative): """ (x ** l_x) * (y ** l_y) * (z ** l_z) / (l_x! * l_y! * l_z!) ** 0.5 """ normalize = 1 / np.sqrt(np.array([1, 1, 1, 1, 1, 1, 1, 2, 2, 2, # 1 / sqrt(lx! ly! lz!) 1, 2, 2, 2, 2, 2, 2, 6, 6, 6])) L_list = [[[0], [0]], # lx = 1, ly = 0, lz = 0; L = 1 [[1], [1]], # lx = 0, ly = 1, lz = 0; L = 1 [[2], [2]], # lx = 0, ly = 0, lz = 1; L = 1 [[0,1], [0,1]], # lx = 1, ly = 1, lz = 0; L = 2 [[0,2], [0,2]], # lx = 1, ly = 0, lz = 1; L = 2 [[1,2], [1,2]], # lx = 0, ly = 1, lz = 1; L = 2 [[0], [3]], # lx = 2, ly = 0, lz = 0; L = 2 [[1], [4]], # lx = 0, ly = 2, lz = 0; L = 2 [[2], [5]], # lx = 0, ly = 0, lz = 2; L = 2 [[0,1,2], [0,1,2]], # lx = 1, ly = 1, lz = 1; L = 3 [[1,2], [1,5]], # lx = 0, ly = 1, lz = 2; L = 3 [[1,2], [4,2]], # lx = 0, ly = 2, lz = 1; L = 3 [[0,2], [0,5]], # lx = 1, ly = 0, lz = 2; L = 3 [[0,1], [0,4]], # lx = 1, ly = 2, lz = 0; L = 3 [[0,1], [3,1]], # lx = 2, ly = 1, lz = 0; L = 3 [[0,2], [3,2]], # lx = 2, ly = 0, lz = 1; L = 3 [[0], [6]], # lx = 3, ly = 0, lz = 0; L = 3 [[1], [7]], # lx = 0, ly = 3, lz = 0; L = 3 [[2], [8]] # lx = 0, ly = 0, lz = 3; L = 3 ] uN = len(unique_js) l = 1 RIJ = np.zeros([len(rij), 9]) dRIJ = np.zeros([len(rij), 9]) if L == 1: l = 4 RIJ[:, :3] = rij if derivative: dRIJ[:, :3] += 1 elif L == 2: l = 10 RIJ[:, :3] = rij RIJ[:, 3:6] = rij*rij if derivative: dRIJ[:, :3] += 1 dRIJ[:, 3:6] = 2*rij elif L == 3: l = 20 RIJ[:, :3] = rij RIJ[:, 3:6] = rij*rij RIJ[:, 6:9] = (rij*rij)*rij if derivative: dRIJ[:, :3] += 1 dRIJ[:, 3:6] = 2*rij dRIJ[:, 6:9] = 3*RIJ[:, 3:6] xyz = np.ones([len(rij), 3, l]) if derivative: dxyz = np.zeros([len(rij), uN, 3, l]) i_dxyz = np.zeros([len(rij), uN, 3, l]) j_dxyz = np.zeros([len(rij), uN, 3, l]) dij_dmlist, i_dij_dmlist, j_dij_dmlist = dij_dm_list(unique_js, ij_list) # [j, uN] for i in range(1, l): xyz[:, L_list[i-1][0], i] = RIJ[:, L_list[i-1][1]] if derivative: # [j, uN], [j, l] -> [j, uN, l] dxyz[:, :, L_list[i-1][0], i] = np.einsum('ij,ik->ijk', dij_dmlist, dRIJ[:, L_list[i-1][1]]) i_dxyz[:, :, L_list[i-1][0], i] = np.einsum('ij,ik->ijk', i_dij_dmlist, dRIJ[:, L_list[i-1][1]]) j_dxyz[:, :, L_list[i-1][0], i] = np.einsum('ij,ik->ijk', j_dij_dmlist, dRIJ[:, L_list[i-1][1]]) result = xyz[:, 0, :] * xyz[:, 1, :] * xyz[:, 2, :] * normalize[:l] # [j, l] if derivative: d_result = np.zeros_like(dxyz) # [j, uN, 3, l] d_result[:, :, 0, :] = np.einsum('ijk,ik->ijk', dxyz[:, :, 0, :], xyz[:, 1, :]*xyz[:, 2, :]) d_result[:, :, 1, :] = np.einsum('ijk,ik->ijk', dxyz[:, :, 1, :], xyz[:, 0, :]*xyz[:, 2, :]) d_result[:, :, 2, :] = np.einsum('ijk,ik->ijk', dxyz[:, :, 2, :], xyz[:, 0, :]*xyz[:, 1, :]) d_result = np.einsum('ijkl,l->ijkl', d_result, normalize[:l]) i_d_result = np.zeros_like(dxyz) # [j, uN, 3, l] i_d_result[:, :, 0, :] = np.einsum('ijk,ik->ijk', i_dxyz[:, :, 0, :], xyz[:, 1, :]*xyz[:, 2, :]) i_d_result[:, :, 1, :] = np.einsum('ijk,ik->ijk', i_dxyz[:, :, 1, :], xyz[:, 0, :]*xyz[:, 2, :]) i_d_result[:, :, 2, :] = np.einsum('ijk,ik->ijk', i_dxyz[:, :, 2, :], xyz[:, 0, :]*xyz[:, 1, :]) i_d_result = np.einsum('ijkl,l->ijkl', i_d_result, normalize[:l]) j_d_result = np.zeros_like(dxyz) # [j, uN, 3, l] j_d_result[:, :, 0, :] = np.einsum('ijk,ik->ijk', j_dxyz[:, :, 0, :], xyz[:, 1, :]*xyz[:, 2, :]) j_d_result[:, :, 1, :] = np.einsum('ijk,ik->ijk', j_dxyz[:, :, 1, :], xyz[:, 0, :]*xyz[:, 2, :]) j_d_result[:, :, 2, :] = np.einsum('ijk,ik->ijk', j_dxyz[:, :, 2, :], xyz[:, 0, :]*xyz[:, 1, :]) j_d_result = np.einsum('ijkl,l->ijkl', j_d_result, normalize[:l]) return result, d_result, i_d_result, j_d_result else: return result, None, None, None
[docs]def dij_dm_list(unique_js, ij_list): """Get the sign of the derivative of x-y-z ** lx-ly-lz. Parameters ---------- uN: int the unique index of the atom that force is acting on. ij_list: list The list of center atom i w.r.t. the neighbors atom j. Returns ------- result: array [j, uN] The signs (+ or -) for dXij_dm (YZ) * dYij_dm (XZ) * dZij_dm (XY) """ uN = len(unique_js) result = np.zeros([len(ij_list), uN]) i_result = np.zeros([len(ij_list), uN]) j_result = np.zeros([len(ij_list), uN]) ijm_list = np.zeros([len(ij_list), 3, uN], dtype=int) ijm_list[:, -1, :] = unique_js ijm_list[:, :2, :] = np.broadcast_to(ij_list[...,None], ij_list.shape+(uN,)) arr = (ijm_list[:, 2, :] == ijm_list[:, 0, :]) result[arr] = -1 i_result[arr] = -1 arr = (ijm_list[:, 2, :] == ijm_list[:, 1, :]) result[arr] = 1 j_result[arr] = 1 arr = (ijm_list[:, 0, :] == ijm_list[:, 1, :]) result[arr] = 0 return result, i_result, j_result # [j, uN]
[docs]def dRij_dRm_norm(Rij, ijm_list): """Calculate the derivative of Rij norm w. r. t. atom m. This term affects only on i and j. Parameters ---------- Rij : array [j, 3] The vector distances of atom i to atom j. ijm_list: array [j, 3] or [j*k, 3] Id list of center atom i, neighbors atom j, and atom m. Returns ------- dRij_m: array [j, 3] The derivative of pair atoms w.r.t. atom m in x, y, z directions. """ dRij_m = np.zeros([len(Rij), 3]) i_dRij_m = np.zeros([len(Rij), 3]) j_dRij_m = np.zeros([len(Rij), 3]) R1ij = np.linalg.norm(Rij, axis=1).reshape([len(Rij),1]) l1 = (ijm_list[:,2]==ijm_list[:,0]) dRij_m[l1, :] = -Rij[l1]/R1ij[l1] i_dRij_m[l1, :] = -Rij[l1]/R1ij[l1] l2 = (ijm_list[:,2]==ijm_list[:,1]) dRij_m[l2, :] = Rij[l2]/R1ij[l2] j_dRij_m[l2, :] = Rij[l2]/R1ij[l2] l3 = (ijm_list[:,0]==ijm_list[:,1]) dRij_m[l3, :] = 0 return dRij_m, i_dRij_m, j_dRij_m
if __name__ == '__main__': import time from ase.build import bulk np.set_printoptions(formatter={'float': '{: 0.4f}'.format}) Rc = 10 parameters1 = {'L': 2, 'eta': [0.036, 0.071], 'Rs': [0]} # Test for stress for a in [5.0]: #, 5.4, 5.8]: si = bulk('Si', 'diamond', a=a, cubic=True) cell = si.get_cell() cell[0,1] += 0.1 si.set_cell(cell) bp = EAD(parameters1, Rc=Rc, derivative=True, stress=True, cutoff='cosine') des = bp.calculate(si) print("G:", des['x']) print("GPrime") print(des['dxdr']) #print(des['rdxdr'][0:8, -1, :, :]) #pprint(np.einsum('ijklm->klm', des['rdxdr']))