use np.testing.assert_almost_equal for array comparing (#1059)

* use `np.testing.assert_almost_equal` for array comparing

* bugfix

* bugfix

* bugfix

* bugfix

* bugfix

* bugfix

* bugfix

Author: njzjz

source/tests/common.py

@@ -258,11 +258,9 @@ def virial_test (inter,
     num_vir = np.transpose(num_vir, [1,0])    
     box3 = dbox[0].reshape([3,3])
     num_vir = np.matmul(num_vir, box3)
-    for ii in range(3):
-        for jj in range(3):
-            testCase.assertAlmostEqual(ana_vir[ii][jj], num_vir[ii][jj],
-                                       places=places, 
-                                       msg = 'virial component %d %d ' % (ii,jj))
+    np.testing.assert_almost_equal(ana_vir, num_vir,
+                                   places, 
+                                   err_msg = 'virial component')
 
 
 

source/tests/test_data_modifier.py

@@ -119,10 +119,9 @@ def _test_fv (self):
             ep, _, __ = dcm.eval(coordp, box, atype, eval_fv = False)
             em, _, __ = dcm.eval(coordm, box, atype, eval_fv = False)
             num_f = -(ep - em) / (2.*hh)
-            for ff in range(nframes):
-                self.assertAlmostEqual(vf[ff,ii], num_f[ff], 
-                                       places = places,
-                                       msg = 'frame %d dof %d does not match' % (ff, ii))
+            np.testing.assert_almost_equal(vf[:,ii].ravel(), num_f.ravel(), 
+                                           places,
+                                           err_msg = 'dof %d does not match' % (ii))
 
         box3 = np.reshape(box, [nframes, 3,3])
         rbox3 = np.linalg.inv(box3)
@@ -150,10 +149,7 @@ def _test_fv (self):
         t_esti = np.matmul(num_deriv, box3)
 
         # print(t_esti, '\n', vv.reshape([-1, 3, 3]))
-        for ff in range(nframes):
-            for ii in range(3):
-                for jj in range(3):                
-                    self.assertAlmostEqual(t_esti[ff][ii][jj], vv[ff,ii*3+jj], 
-                                           places = places,
-                                           msg = "frame %d virial component [%d,%d] failed" % (ff, ii, jj))
+        np.testing.assert_almost_equal(t_esti.ravel(), vv.ravel(), 
+                                       places,
+                                       err_msg = "virial component failed")
 

source/tests/test_data_modifier_shuffle.py

@@ -185,11 +185,9 @@ def test_z_dipole(self):
         dv01 = dv01.reshape([self.nframes, -1])
         dv1 = dv1.reshape([self.nframes, -1])
 
-        for ii in range(self.nframes):
-            for jj in range(self.nsel):
-                self.assertAlmostEqual(
-                    dv01[ii][jj], dv1[ii][jj], 
-                    msg = "dipole [%d,%d] dose not match" % (ii, jj))
+        np.testing.assert_almost_equal(
+                    dv01, dv1, 
+                    err_msg = "dipole dose not match")
 
 
     def test_modify(self):
@@ -202,18 +200,12 @@ def test_modify(self):
         ve1, vf1, vv1 = dcm.eval(self.coords1, self.box1, self.atom_types1)
         vf01 = vf0[:,self.idx_map, :]
 
-        for ii in range(self.nframes):
-            self.assertAlmostEqual(ve0[ii], ve1[ii], 
-                                   msg = 'energy %d should match' % ii)
-        for ii in range(self.nframes):
-            for jj in range(9):
-                self.assertAlmostEqual(vv0[ii][jj], vv1[ii][jj], 
-                                       msg = 'virial [%d,%d] should match' % (ii,jj))
-        for ii in range(self.nframes):
-            for jj in range(self.natoms):
-                for dd in range(3):
-                    self.assertAlmostEqual(
-                        vf01[ii][jj][dd], vf1[ii][jj][dd], 
-                        msg = "force [%d,%d,%d] dose not match" % (ii,jj,dd))
+        np.testing.assert_almost_equal(ve0, ve1, 
+                                       err_msg = 'energy should match')
+        np.testing.assert_almost_equal(vv0, vv1, 
+                                       err_msg = 'virial should match')
+        np.testing.assert_almost_equal(
+                        vf01, vf1, 
+                        err_msg = "force dose not match")
 
 

source/tests/test_deepdipole.py

@@ -44,8 +44,7 @@ def test_1frame_atm(self):
         nsel = 2
         self.assertEqual(dd.shape, (nframes,nsel,3))
         # check values
-        for ii in range(dd.size):
-            self.assertAlmostEqual(dd.reshape([-1])[ii], self.expected_d.reshape([-1])[ii], places = default_places)
+        np.testing.assert_almost_equal(dd.ravel(), self.expected_d, default_places)
 
     def test_2frame_atm(self):
         coords2 = np.concatenate((self.coords, self.coords))
@@ -58,8 +57,7 @@ def test_2frame_atm(self):
         self.assertEqual(dd.shape, (nframes,nsel,3))
         # check values
         expected_d = np.concatenate((self.expected_d, self.expected_d))
-        for ii in range(dd.size):
-            self.assertAlmostEqual(dd.reshape([-1])[ii], expected_d.reshape([-1])[ii], places = default_places)
+        np.testing.assert_almost_equal(dd.ravel(), expected_d, default_places)
 
 
 class TestDeepDipoleNoPBC(unittest.TestCase) :
@@ -87,8 +85,7 @@ def test_1frame_atm(self):
         nsel = 2
         self.assertEqual(dd.shape, (nframes,nsel,3))
         # check values
-        for ii in range(dd.size):
-            self.assertAlmostEqual(dd.reshape([-1])[ii], self.expected_d.reshape([-1])[ii], places = default_places)
+        np.testing.assert_almost_equal(dd.ravel(), self.expected_d, default_places)
 
     def test_1frame_atm_large_box(self):
         dd = self.dp.eval(self.coords, self.box, self.atype)
@@ -98,8 +95,7 @@ def test_1frame_atm_large_box(self):
         nsel = 2
         self.assertEqual(dd.shape, (nframes,nsel,3))
         # check values
-        for ii in range(dd.size):
-            self.assertAlmostEqual(dd.reshape([-1])[ii], self.expected_d.reshape([-1])[ii], places = default_places)
+        np.testing.assert_almost_equal(dd.ravel(), self.expected_d, default_places)
 
 
 @unittest.skipIf(parse_version(tf.__version__) < parse_version("1.15"), 
@@ -138,8 +134,7 @@ def test_1frame_old(self):
         nframes = 1
         self.assertEqual(gt.shape, (nframes,self.nout))
         # check values
-        for ii in range(gt.size):
-            self.assertAlmostEqual(gt.reshape([-1])[ii], self.expected_gt.reshape([-1])[ii], places = default_places)
+        np.testing.assert_almost_equal(gt.ravel(), self.expected_gt, default_places)
 
     def test_1frame_old_atm(self):
         at = self.dp.eval(self.coords, self.box, self.atype)
@@ -149,8 +144,7 @@ def test_1frame_old_atm(self):
         nsel = 2
         self.assertEqual(at.shape, (nframes,nsel,self.nout))
         # check values
-        for ii in range(at.size):
-            self.assertAlmostEqual(at.reshape([-1])[ii], self.expected_t.reshape([-1])[ii], places = default_places)
+        np.testing.assert_almost_equal(at.ravel(), self.expected_t, default_places)
 
     def test_2frame_old_atm(self):
         coords2 = np.concatenate((self.coords, self.coords))
@@ -163,8 +157,7 @@ def test_2frame_old_atm(self):
         self.assertEqual(at.shape, (nframes,nsel,self.nout))
         # check values
         expected_d = np.concatenate((self.expected_t, self.expected_t))
-        for ii in range(at.size):
-            self.assertAlmostEqual(at.reshape([-1])[ii], expected_d.reshape([-1])[ii], places = default_places)
+        np.testing.assert_almost_equal(at.ravel(), expected_d, default_places)
 
     def test_1frame_full(self):
         gt, ff, vv = self.dp.eval_full(self.coords, self.box, self.atype, atomic = False)
@@ -175,12 +168,9 @@ def test_1frame_full(self):
         self.assertEqual(ff.shape, (nframes,self.nout,natoms,3))
         self.assertEqual(vv.shape, (nframes,self.nout,9))
         # check values
-        for ii in range(ff.size):
-            self.assertAlmostEqual(ff.reshape([-1])[ii], self.expected_f.reshape([-1])[ii], places = default_places)
-        for ii in range(gt.size):
-            self.assertAlmostEqual(gt.reshape([-1])[ii], self.expected_gt.reshape([-1])[ii], places = default_places)
-        for ii in range(vv.size):
-            self.assertAlmostEqual(vv.reshape([-1])[ii], self.expected_gv.reshape([-1])[ii], places = default_places)
+        np.testing.assert_almost_equal(ff.ravel(), self.expected_f, default_places)
+        np.testing.assert_almost_equal(gt.ravel(), self.expected_gt, default_places)
+        np.testing.assert_almost_equal(vv.ravel(), self.expected_gv, default_places)
 
     def test_1frame_full_atm(self):
         gt, ff, vv, at, av = self.dp.eval_full(self.coords, self.box, self.atype, atomic = True)

source/tests/test_deepmd_data.py

@@ -49,21 +49,15 @@ def test_load_set_1(self) :
              .add('value_1', 1, atomic=True, must=True, type_sel = [0])
         data = dd._load_set(os.path.join(self.data_name, 'set.foo'))
         self.assertEqual(data['value_1'].shape, (self.nframes, 2))
-        for ii in range(self.nframes):
-            for jj in range(2):
-                self.assertAlmostEqual(data['value_1'][ii][jj],
-                                       self.value_1[ii][jj])
+        np.testing.assert_almost_equal(data['value_1'], self.value_1)
 
 
     def test_load_set_2(self) :
         dd = DeepmdData(self.data_name)\
              .add('value_2', 1, atomic=True, must=True, type_sel = [1])
         data = dd._load_set(os.path.join(self.data_name, 'set.foo'))
         self.assertEqual(data['value_2'].shape, (self.nframes, 4))
-        for ii in range(self.nframes):
-            for jj in range(4):
-                self.assertAlmostEqual(data['value_2'][ii][jj],
-                                       self.value_2[ii][jj])                
+        np.testing.assert_almost_equal(data['value_2'], self.value_2)          
 
 
 class TestData (unittest.TestCase) :
@@ -217,8 +211,7 @@ def test_avg(self) :
              .add('test_frame', 5, atomic=False, must=True)
         favg = dd.avg('test_frame')
         fcmp = np.average(np.concatenate((self.test_frame, self.test_frame_bar), axis = 0), axis = 0)
-        for ii in range(favg.size) :
-            self.assertAlmostEqual((favg[ii]), (fcmp[ii]), places = places)
+        np.testing.assert_almost_equal(favg, fcmp, places)
 
     def test_check_batch_size(self) :
         dd = DeepmdData(self.data_name)
@@ -263,8 +256,4 @@ def test_get_nbatch(self):
         self.assertEqual(nb, 2)
 
     def _comp_np_mat2(self, first, second) :
-        for ii in range(first.shape[0]) :
-            for jj in range(first.shape[1]) :
-                self.assertAlmostEqual(first[ii][jj], second[ii][jj], 
-                                       msg = 'item [%d][%d] does not match' % (ii,jj), 
-                                       places = places)
+        np.testing.assert_almost_equal(first, second, places)

source/tests/test_deeppolar.py

@@ -44,8 +44,7 @@ def test_1frame_atm(self):
         nsel = 2
         self.assertEqual(dd.shape, (nframes,nsel,9))
         # check values
-        for ii in range(dd.size):
-            self.assertAlmostEqual(dd.reshape([-1])[ii], self.expected_d.reshape([-1])[ii], places = default_places)
+        np.testing.assert_almost_equal(dd.ravel(), self.expected_d, default_places)
 
     def test_2frame_atm(self):
         coords2 = np.concatenate((self.coords, self.coords))
@@ -58,8 +57,8 @@ def test_2frame_atm(self):
         self.assertEqual(dd.shape, (nframes,nsel,9))
         # check values
         expected_d = np.concatenate((self.expected_d, self.expected_d))
-        for ii in range(dd.size):
-            self.assertAlmostEqual(dd.reshape([-1])[ii], expected_d.reshape([-1])[ii], places = default_places)
+        np.testing.assert_almost_equal(dd.ravel(), expected_d, default_places)
+
 
 
 class TestDeepPolarNoPBC(unittest.TestCase) :
@@ -87,8 +86,7 @@ def test_1frame_atm(self):
         nsel = 2
         self.assertEqual(dd.shape, (nframes,nsel,9))
         # check values
-        for ii in range(dd.size):
-            self.assertAlmostEqual(dd.reshape([-1])[ii], self.expected_d.reshape([-1])[ii], places = default_places)
+        np.testing.assert_almost_equal(dd.ravel(), self.expected_d, default_places)
 
     def test_1frame_atm_large_box(self):
         dd = self.dp.eval(self.coords, self.box, self.atype)
@@ -98,8 +96,7 @@ def test_1frame_atm_large_box(self):
         nsel = 2
         self.assertEqual(dd.shape, (nframes,nsel,9))
         # check values
-        for ii in range(dd.size):
-            self.assertAlmostEqual(dd.reshape([-1])[ii], self.expected_d.reshape([-1])[ii], places = default_places)
+        np.testing.assert_almost_equal(dd.ravel(), self.expected_d, default_places)
 
 
 @unittest.skipIf(parse_version(tf.__version__) < parse_version("1.15"), 
@@ -138,8 +135,7 @@ def test_1frame_old(self):
         nframes = 1
         self.assertEqual(gt.shape, (nframes,self.nout))
         # check values
-        for ii in range(gt.size):
-            self.assertAlmostEqual(gt.reshape([-1])[ii], self.expected_gt.reshape([-1])[ii], places = default_places)
+        np.testing.assert_almost_equal(gt.ravel(), self.expected_gt, default_places)
 
     def test_1frame_old_atm(self):
         at = self.dp.eval(self.coords, self.box, self.atype)
@@ -149,8 +145,7 @@ def test_1frame_old_atm(self):
         nsel = 2
         self.assertEqual(at.shape, (nframes,nsel,self.nout))
         # check values
-        for ii in range(at.size):
-            self.assertAlmostEqual(at.reshape([-1])[ii], self.expected_t.reshape([-1])[ii], places = default_places)
+        np.testing.assert_almost_equal(at.ravel(), self.expected_t, default_places)
 
     def test_2frame_old_atm(self):
         coords2 = np.concatenate((self.coords, self.coords))
@@ -163,8 +158,7 @@ def test_2frame_old_atm(self):
         self.assertEqual(at.shape, (nframes,nsel,self.nout))
         # check values
         expected_d = np.concatenate((self.expected_t, self.expected_t))
-        for ii in range(at.size):
-            self.assertAlmostEqual(at.reshape([-1])[ii], expected_d.reshape([-1])[ii], places = default_places)
+        np.testing.assert_almost_equal(at.ravel(), expected_d, default_places)
 
     def test_1frame_full(self):
         gt, ff, vv = self.dp.eval_full(self.coords, self.box, self.atype, atomic = False)
@@ -175,12 +169,9 @@ def test_1frame_full(self):
         self.assertEqual(ff.shape, (nframes,self.nout,natoms,3))
         self.assertEqual(vv.shape, (nframes,self.nout,9))
         # check values
-        for ii in range(ff.size):
-            self.assertAlmostEqual(ff.reshape([-1])[ii], self.expected_f.reshape([-1])[ii], places = default_places)
-        for ii in range(gt.size):
-            self.assertAlmostEqual(gt.reshape([-1])[ii], self.expected_gt.reshape([-1])[ii], places = default_places)
-        for ii in range(vv.size):
-            self.assertAlmostEqual(vv.reshape([-1])[ii], self.expected_gv.reshape([-1])[ii], places = default_places)
+        np.testing.assert_almost_equal(ff.ravel(), self.expected_f, default_places)
+        np.testing.assert_almost_equal(gt.ravel(), self.expected_gt, default_places)
+        np.testing.assert_almost_equal(vv.ravel(), self.expected_gv, default_places)
 
     def test_1frame_full_atm(self):
         gt, ff, vv, at, av = self.dp.eval_full(self.coords, self.box, self.atype, atomic = True)