use self.date_idx for selecting train_idx and test_idx

stefan-jansen · Apr 9, 2021 · a094633 · a094633
1 parent 7ac6764
commit a094633
Showing 1 changed file with 3 additions and 3 deletions.
diff --git a/utils.py b/utils.py
@@ -50,9 +50,9 @@ def split(self, X, y=None, groups=None):
         for train_start, train_end, test_start, test_end in split_idx:
 
             train_idx = dates[(dates[self.date_idx] > days[train_start])
-                              & (dates.date <= days[train_end])].index
-            test_idx = dates[(dates.date > days[test_start])
-                             & (dates.date <= days[test_end])].index
+                              & (dates[self.date_idx] <= days[train_end])].index
+            test_idx = dates[(dates[self.date_idx] > days[test_start])
+                             & (dates[self.date_idx] <= days[test_end])].index
             if self.shuffle:
                 np.random.shuffle(list(train_idx))
             yield train_idx.to_numpy(), test_idx.to_numpy()