From 534edc5aefe43a8966c164208118f7b24e5c8449 Mon Sep 17 00:00:00 2001 From: VScoldness Date: Fri, 5 Jan 2024 17:17:25 -0700 Subject: [PATCH] Update neuralnetwork class - move calculation of dxdr from calculate_loss into normalize function - modified the calculation of dxdr in calculate_loss and evaluate functions --- pyxtal_ff/models/neuralnetwork.py | 28 ++++++++++++++++------------ 1 file changed, 16 insertions(+), 12 deletions(-) diff --git a/pyxtal_ff/models/neuralnetwork.py b/pyxtal_ff/models/neuralnetwork.py index 16ffc53..7232b6e 100644 --- a/pyxtal_ff/models/neuralnetwork.py +++ b/pyxtal_ff/models/neuralnetwork.py @@ -308,13 +308,7 @@ def evaluate(self, data, figname): _Energy += _e if self.force_coefficient: dedx = torch.autograd.grad(_e, _x)[0] - _dxdr = value['dxdr'][element] - - tmp = np.zeros([len(value['x'][element]), n_atoms, value['x'][element].shape[1], 3]) - for _m in range(n_atoms): - rows = np.where(value['seq'][element][:,1]==_m)[0] - tmp[value['seq'][element][rows, 0], _m, :, :] += _dxdr[rows, :, :] - _Force -= torch.einsum("ik, ijkl->jl", dedx, torch.from_numpy(tmp)) + _Force -= torch.einsum("ik, ijkl->jl", dedx, value['dxdr'][element]) if self.stress_coefficient and (data2['group'] in self.stress_group): if self.force_coefficient is None: @@ -495,11 +489,7 @@ def calculate_loss(self, models, batch): if self.force_coefficient: dedx[element] = torch.autograd.grad(_energy, _x, create_graph=True)[0] - tmp = np.zeros([len(x[element]), n_atoms, x[element].shape[1], 3]) - for _m in range(n_atoms): - rows = np.where(seq[element][:,1]==_m)[0] - tmp[seq[element][rows, 0], _m, :, :] += dxdr[element][rows, :, :] - _force -= torch.einsum("ik, ijkl->jl", dedx[element], torch.from_numpy(tmp)) + _force -= torch.einsum("ik, ijkl->jl", dedx[element], dxdr[element]) if self.stress_coefficient and (group in self.stress_group): if self.force_coefficient is None: @@ -1019,6 +1009,20 @@ def normalize(self, data, drange, unit, norm=[0., 1.]): #d['seq'][element][:, 0] -= torch.min(d['seq'][element][:, 0]) #adjust the initial position d['seq'][element][:, 0] -= np.min(d['seq'][element][:, 0]) #adjust the initial position + x, dxdr, seq = d['x'], d['dxdr'], d['seq'] + n_atoms = sum(len(value) for value in x.values()) + new_dxdr = {} + for element in x.keys(): + if x[element].nelement() > 0: + tmp = np.zeros([len(x[element]), n_atoms, x[element].shape[1], 3]) + for _m in range(n_atoms): + rows = np.where(seq[element][:,1]==_m)[0] + tmp[seq[element][rows, 0], _m, :, :] += dxdr[element][rows, :, :] + new_dxdr[element] = torch.from_numpy(tmp) + else: + new_dxdr[element] = torch.empty((30,3), dtype=torch.float) + d['dxdr'] = new_dxdr + db2[str(i)] = d db1.close()