Merge pull request #79 from njzjz/shuffle

add shuffle method to shuffle data

Author: amcadmus

dpdata/system.py

@@ -773,6 +773,14 @@ def nopbc(self):
             return True
         return False
 
+    def shuffle(self):
+        """Shuffle frames randomly."""
+        idx = np.random.permutation(self.get_nframes())
+        for ii in ['cells', 'coords']:
+            self.data[ii] = self.data[ii][idx]
+        return idx
+
+
 def get_cell_perturb_matrix(cell_pert_fraction):
     if cell_pert_fraction<0:
         raise RuntimeError('cell_pert_fraction can not be negative')
@@ -1131,6 +1139,14 @@ def sort_atom_types(self):
             if ii in self.data:
                 self.data[ii] = self.data[ii][:, idx]
 
+    def shuffle(self):
+        """Also shuffle labeled data e.g. energies and forces."""
+        idx = System.shuffle(self)
+        for ii in ['energies', 'forces', 'virials', 'atom_pref']:
+            if ii in self.data:
+                self.data[ii] = self.data[ii][idx]
+        return idx
+
 
 class MultiSystems:
     '''A set containing several systems.'''

tests/test_qe_pw_scf.py

@@ -38,6 +38,7 @@ def test_cell(self) :
         for ii in range(cell.shape[0]) :
             for jj in range(cell.shape[1]) :
                 self.assertAlmostEqual(self.system_h2o.data['cells'][0][ii][jj], cell[ii][jj])
+        fp.close()
 
 
     def test_coord(self) :

tests/test_shuffle.py

@@ -0,0 +1,21 @@
+import unittest
+from context import dpdata
+from comp_sys import CompLabeledSys, IsPBC
+
+class TestDeepmdLoadRaw(unittest.TestCase, CompLabeledSys, IsPBC):
+    def setUp (self) :
+        original_system = dpdata.LabeledSystem('poscars/OUTCAR.h2o.md', 
+                                             fmt = 'vasp/outcar')
+        original_system += original_system
+        original_system += original_system
+        original_system += original_system
+        self.system_1 = dpdata.LabeledSystem()
+        self.system_2 = original_system.copy()
+        idx = self.system_2.shuffle()
+        for ii in idx:
+            self.system_1.append(original_system.sub_system(ii))
+
+        self.places = 6
+        self.e_places = 6
+        self.f_places = 6
+        self.v_places = 6