From 0b93d91c7b5b27413f68b717a6fa57a5b4435c51 Mon Sep 17 00:00:00 2001
From: Jinzhe Zeng <jinzhe.zeng@rutgers.edu>
Date: Tue, 14 Jan 2020 14:32:24 -0500
Subject: [PATCH] add shuffle method to shuffle data

---
 dpdata/system.py        | 16 ++++++++++++++++
 tests/test_qe_pw_scf.py |  1 +
 tests/test_shuffle.py   | 21 +++++++++++++++++++++
 3 files changed, 38 insertions(+)
 create mode 100644 tests/test_shuffle.py

diff --git a/dpdata/system.py b/dpdata/system.py
index 12cb4692e..9ba6ad840 100644
--- a/dpdata/system.py
+++ b/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.'''
diff --git a/tests/test_qe_pw_scf.py b/tests/test_qe_pw_scf.py
index eb1c96bb7..391e763eb 100644
--- a/tests/test_qe_pw_scf.py
+++ b/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) :
diff --git a/tests/test_shuffle.py b/tests/test_shuffle.py
new file mode 100644
index 000000000..aa2d7a7e3
--- /dev/null
+++ b/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