fix data split issue (#149)

choderalab · Apr 20, 2023 · 821c5f2 · 821c5f2
1 parent b2b3f20
commit 821c5f2
Showing 1 changed file with 6 additions and 2 deletions.
diff --git a/espaloma/data/dataset.py b/espaloma/data/dataset.py
@@ -186,10 +186,14 @@ def split(self, partition):
 
         """
         n_data = len(self)
-        partition = [int(n_data * x / sum(partition)) for x in partition]
+        p_sizes = []
+        for i, _partition in enumerate(partition):
+            p_size = int((n_data - sum(p_sizes)) * _partition / sum(partition[i:]))
+            p_sizes.append(p_size)
+        assert sum(p_sizes) == n_data, f"{p_sizes}, {sum(p_sizes)}"
         ds = []
         idx = 0
-        for p_size in partition:
+        for p_size in p_sizes:
             ds.append(self[idx : idx + p_size])
             idx += p_size