fix bug in dp test

Author: Han Wang

source/train/Data.py

@@ -359,6 +359,14 @@ def __init__ (self,
             self.has_fparam = 0
         else :
             self.has_fparam = -1
+        # check aparam
+        has_aparam = [ os.path.isfile(os.path.join(ii, 'aparam.npy')) for ii in self.dirs ]
+        if any(has_aparam) and (not all(has_aparam)) :
+            raise RuntimeError("system %s: if any set has frame parameter, then all sets should have frame parameter" % sys_path)
+        if all(has_aparam) :
+            self.has_aparam = 0
+        else :
+            self.has_aparam = -1
         # energy norm
         self.eavg = self.stats_energy()
         # load sets
@@ -463,6 +471,12 @@ def load_set(self, set_name, shuffle = True):
                 self.has_fparam = data["fparam"].shape[1]
             else :
                 assert self.has_fparam == data["fparam"].shape[1]
+        if self.has_aparam >= 0:
+            data["aparam"] = self.load_data(set_name, "aparam", [nframe, -1])
+            if self.has_aparam == 0 :
+                self.has_aparam = data["aparam"].shape[1] // (ncoord//3)
+            else :
+                assert self.has_aparam == data["aparam"].shape[1] // (ncoord//3)
         data["prop_c"] = np.zeros(5)
         data["prop_c"][0], data["energy"], data["prop_c"][3], data["atom_ener"] \
             = self.load_energy (set_name, nframe, ncoord // 3, "energy", "atom_ener")
@@ -573,3 +587,6 @@ def get_ener (self) :
     def numb_fparam(self) :
         return self.has_fparam
 
+    def numb_aparam(self) :
+        return self.has_aparam
+

source/train/test.py

@@ -42,7 +42,15 @@ def test_ener (args) :
     coord = test_data["coord"][:numb_test].reshape([numb_test, -1])
     box = test_data["box"][:numb_test]
     atype = test_data["type"][0]
-    energy, force, virial, ae, av = dp.eval(coord, box, atype, fparam = (test_data["fparam"] if "fparam" in test_data else None), atomic = True)
+    if dp.get_dim_fparam() > 0:
+        fparam = test_data["fparam"][:numb_test] 
+    else :
+        fparam = None
+    if dp.get_dim_aparam() > 0:
+        aparam = test_data["aparam"][:numb_test] 
+    else :
+        aparam = None
+    energy, force, virial, ae, av = dp.eval(coord, box, atype, fparam = fparam, aparam = aparam, atomic = True)
     energy = energy.reshape([numb_test,1])
     force = force.reshape([numb_test,-1])
     virial = virial.reshape([numb_test,9])