fix bug #165

Author: Han Wang

source/train/test.py

@@ -36,13 +36,19 @@ def test_ener (args) :
 
     dp = DeepPot(args.model)
     data = DeepmdData(args.system, args.set_prefix, shuffle_test = args.shuffle_test, type_map = dp.get_type_map())
+    data.add('energy', 1, atomic=False, must=False, high_prec=True)
+    data.add('force',  3, atomic=True,  must=False, high_prec=False)
+    data.add('virial', 9, atomic=False, must=False, high_prec=False)
+    if dp.get_dim_fparam() > 0:
+        data.add('fparam', dp.get_dim_fparam(), atomic=False, must=True, high_prec=False)
+    if dp.get_dim_aparam() > 0:
+        data.add('aparam', dp.get_dim_aparam(), atomic=True,  must=True, high_prec=False)
 
     test_data = data.get_test ()
-    numb_test = args.numb_test
     natoms = len(test_data["type"][0])
     nframes = test_data["box"].shape[0]
+    numb_test = args.numb_test
     numb_test = min(nframes, numb_test)
-
     coord = test_data["coord"][:numb_test].reshape([numb_test, -1])
     box = test_data["box"][:numb_test]
     atype = test_data["type"][0]
@@ -54,6 +60,7 @@ def test_ener (args) :
         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])