The return of read_k_cylinder and line_to_state has been changed.

deGrootLab · Aug 6, 2023 · 07eedd4 · 07eedd4
1 parent 4c7a0b5
commit 07eedd4
Show file tree

Hide file tree

Showing 3 changed files with 8 additions and 5 deletions.
diff --git a/Sfilter/util/MSM.py b/Sfilter/util/MSM.py
@@ -103,7 +103,7 @@ def __init__(self, file_list, start=0, end=None, step=1, method="K_priority"):
             pass
         elif isinstance(file_list, str):
             file_list = [file_list]
-        for traj, meta_data in [read_k_cylinder(file, method) for file in file_list]:
+        for traj, meta_data, K_occupency, W_occupency in [read_k_cylinder(file, method) for file in file_list]:
             self.state_str.append(traj[start:end:step])
             self.time_step.append(meta_data["time_step"] * step)
         flattened = [num for sublist in self.state_str for num in sublist]

diff --git a/test/test_MSM.py b/test/test_MSM.py
@@ -5,13 +5,13 @@
 
 class MyTestCase(unittest.TestCase):
     def test_MSM_read_k_cylinder(self):
-        s_list, time_step = MSM.read_k_cylinder("04-output_wrapper/C_0.75_2ps/05-2ps/00/analysis/04-state-code/k_cylinder.log")
+        s_list, time_step, K_occupency, W_occupency = MSM.read_k_cylinder("04-output_wrapper/C_0.75_2ps/05-2ps/00/analysis/04-state-code/k_cylinder.log")
         self.assertListEqual(s_list[:3], ["WK0KKW", "WK0KKW", "WK0KKW"])
 
 
 
     def test_read_state_file_co_occupy(self):
-        s_list, time_step = MSM.read_k_cylinder("04-output_wrapper/C_0.75_2ps/05-2ps/00/analysis/04-state-code/k_cylinder.log", method="Co-occupy")
+        s_list, time_step, K_occupency, W_occupency = MSM.read_k_cylinder("04-output_wrapper/C_0.75_2ps/05-2ps/00/analysis/04-state-code/k_cylinder.log", method="Co-occupy")
         self.assertListEqual(s_list[:3], ["WK0KKW", "WK0KKW", "WK0KKW"])
         self.assertListEqual(s_list[405:408], ["WK0KKC", "WK0KKC", "WKK0KW"])
 

diff --git a/test/test_output_wrapper.py b/test/test_output_wrapper.py
@@ -11,8 +11,11 @@ def test_line_to_state(self):
                  " POT : 5961 , Wat : ,\n",
                  " POT : 5960 , Wat : ,\n",
                  " POT : 5959 , Wat : 29864 38630 40004 40109 42965 43178 44273 44456 44522 46910 47849 48983 55532 60080 60440 ,\n",]
-        for l, answer in zip(lines, ["W", "K", "0", "K", "K", "C"]):
-            self.assertEqual(Sfilter.util.output_wrapper.line_to_state(l), answer)
+        for l, state_str_ref, n_pot_ref, n_wat_ref  in zip(lines, ["W", "K", "0", "K", "K", "C"], [0, 1, 0, 1, 1, 1], [2, 0, 0, 0, 0, 15]):
+            state_str, n_pot, n_wat = Sfilter.util.output_wrapper.line_to_state(l)
+            self.assertEqual(state_str, state_str_ref)
+            self.assertEqual(n_pot, n_pot_ref)
+            self.assertEqual(n_wat, n_wat_ref)
 
     def test_Perm_event_output_read_file(self):
         files = ["03-longest_common_sequence/03-state-code/POT_perm_event.out",