Bugfix

pyxtal_ff/__init__.py

@@ -53,7 +53,7 @@ def __init__(self, descriptors=None, model=None, logo=True):
                     {'lmax': 3}
                 + SO3
                     {'nmax': 1, 'lmax': 3, 'alpha': 2.0}
-                + SNAP
+                + SNAP (not in parameters, but in descriptors dict directly)
                     {'weights': {'Si': 1.0, 'O': 2.0},
                      'Rc': {'Si': 4.0, 'O': 5.0}
             - zbl: dict

pyxtal_ff/models/neuralnetwork.py

@@ -499,7 +499,9 @@ def calculate_loss(self, models, batch):
                         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)) 
+                        tmp_torch = torch.from_numpy(tmp)
+                        tmp_torch = tmp_torch.to(self.device)
+                        _force -= torch.einsum("ik, ijkl->jl", dedx[element], tmp_torch)
 
                     if self.stress_coefficient and (group in self.stress_group):
                         if self.force_coefficient is None:

pyxtal_ff/models/polynomialregression.py

@@ -81,7 +81,7 @@ def train(self, TrainData, optimizer):
             self.d_max = db['0']['x'].shape[1]
         else:
             # d_max has to be less or equal than total descriptors.
-            assert self.d_max <= len(db['0']['x'].shape[1]),\
+            assert self.d_max <= db['0']['x'].shape[1],\
                     "d_max is larger than total descriptors."
 
         if self.stress_coefficient and (self.stress_group is None):

pyxtal_ff/utilities/__init__.py

@@ -133,7 +133,7 @@ def add(self, function, data):
 
         if _cpu == 1:
             for i, index in enumerate(lists):
-                d = self.compute(function, data[index])
+                d = compute_descriptor(function, data[index], base_potential=self.base_potential)
                 self.append(d)
                 self.length += 1
                 print('\r{:4d} out of {:4d}'.format(i+1, len(lists)), flush=True, end='')
@@ -143,7 +143,7 @@ def add(self, function, data):
             _data = [data[item] for item in lists]
             failed_ids = []
             with Pool(_cpu) as p:
-                func = partial(self.compute, function)
+                func = partial(compute_descriptor, function, base_potential=self.base_potential)
                 for i, d in enumerate(p.imap_unordered(func, _data)):
                     try:
                         self.append(d)
@@ -157,13 +157,14 @@ def add(self, function, data):
 
             # compute the missing structures in parallel calcs
             for id in failed_ids: 
-                d = self.compute(function, _data[id])
+                d = compute_descriptor(function, _data[id], base_potential=self.base_potential)
                 self.append(d)
                 self.length += 1
 
         print(f"\nSaving descriptor-feature data to {self.name}.dat\n")
 
 
+    '''
     def compute(self, function, data):
         """ Compute descriptor for one structure to the database. """
 
@@ -252,9 +253,10 @@ def compute(self, function, data):
         d['group'] = data['group']
 
         return d
+    '''
 
 
-def compute_descriptor(function, structure):
+def compute_descriptor(function, data, base_potential=None):
     """ Compute descriptor for one structure. """
 
     if function['type'] in ['BehlerParrinello', 'ACSF']:
@@ -263,15 +265,15 @@ def compute_descriptor(function, structure):
                  function['Rc'], 
                  function['force'],
                  function['stress'],
-                 function['cutoff'], False).calculate(structure)
+                 function['cutoff'], False).calculate(data['structure'])
 
     elif function['type'] in ['wACSF', 'wacsf']:
         from pyxtal_ff.descriptors.ACSF import ACSF
         d = ACSF(function['parameters'],
                  function['Rc'], 
                  function['force'],
                  function['stress'],
-                 function['cutoff'], True).calculate(structure)
+                 function['cutoff'], True).calculate(data['structure'])
 
     elif function['type'] in ['SO4', 'Bispectrum', 'bispectrum']:
         from pyxtal_ff.descriptors.SO4 import SO4_Bispectrum
@@ -280,22 +282,22 @@ def compute_descriptor(function, structure):
                            derivative=True,
                            stress=True,
                            normalize_U=function['parameters']['normalize_U'],
-                           cutoff_function=function['cutoff']).calculate(structure)
+                           cutoff_function=function['cutoff']).calculate(data['structure'])
 
     elif function['type'] in ['SO3', 'SOAP', 'soap']:
         from pyxtal_ff.descriptors.SO3 import SO3
         d = SO3(function['parameters']['nmax'],
                 function['parameters']['lmax'],
                 function['Rc'],
                 derivative=True,
-                stress=True).calculate(structure)
+                stress=True).calculate(data['structure'])
 
     elif function['type'] in ['EAD', 'ead']:
             from pyxtal_ff.descriptors.EAD import EAD
             d = EAD(function['parameters'],
                      function['Rc'],
                      True, True,
-                     function['cutoff']).calculate(structure)
+                     function['cutoff']).calculate(data['structure'])
 
     elif function['type'] in ['SNAP', 'snap']:
         from pyxtal_ff.descriptors.SNAP import SO4_Bispectrum
@@ -305,7 +307,7 @@ def compute_descriptor(function, structure):
                            derivative=True,
                            stress=True,
                            normalize_U=function['parameters']['normalize_U'],
-                           cutoff_function=function['cutoff']).calculate(structure)
+                           cutoff_function=function['cutoff']).calculate(data['structure'])
 
     else:
         msg = f"{function['type']} is not implemented"
@@ -320,6 +322,19 @@ def compute_descriptor(function, structure):
             rdxdr[_m, :, :, :] += np.einsum('ijkl->jkl', d['rdxdr'][ids, :, :, :])
         d['rdxdr'] = rdxdr.reshape([N, L, 9])[:, :, [0, 4, 8, 1, 2, 5]]
 
+    if base_potential:
+        base_d = base_potential.calculate(data['structure'])
+    else:
+        base_d = {'energy': 0., 'force': 0., 'stress': 0.}
+
+    d['energy'] = np.asarray(data['energy'] - base_d['energy'])
+    d['force'] = np.asarray(data['force']) - base_d['force']
+    if data['stress'] is not None:
+        d['stress'] = np.asarray(data['stress']) - base_d['stress'] / units.GPa
+    else:
+        d['stress'] = data['stress']
+    d['group'] = data['group']
+
     return d