LeavePAdjacentOut, LeaveMapGroupsOut

leavepout_extension.py

# author: p4perf4ce

class LeavePAdjacentOut:
    def __init__(self, p: int = 1, stride: int = 1):
      """
      LeavePAdjacentOut
      
      Implement Leave P-Adjacent Out, a sliding window style LeavePOut instead of the combinatorial implemented in scikit-learn.
      """
        self.p = p
        self.stride = stride

    def split(self, x):
        for start_ind in range(0, len(x), self.stride):
            ind = np.arange(len(x))
            mask = np.ones(len(x), dtype=bool)
            end_ind = start_ind + self.p
            if end_ind > len(x): break
            test_ind = ind[np.arange(start_ind, end_ind)]
            mask[np.arange(start_ind, end_ind)] = False
            train_ind = ind[mask]
            yield train_ind, test_ind

class LeaveMapGroupsOut:
    def __init__(self, class_map, stride: int = 1):
      """
      LeaveMapGroupsOut
      
      Implement a multi-class Leave P-Adjacent Out.
      
      Example:
      >>> lpgo = LeaveMapGroupOut({0: 2, 1: 2})
      >>> a = np.array([0, 0, 1, 1])
      >>> lpgo.split(a)
      >>> [([1, 3], [0, 2]),
      >>> ([1, 2], [0, 3]),
      >>> ...
      """
        self.class_map = class_map
        self.stride = stride
        
    def split(self, y, verbose: bool = False):
        class_map = self.class_map
        lpo_dict = {}
        y_dict = {}
        y_ind_arr = np.arange(len(y))
        classes = list(class_map.keys())
        for c, n in class_map.items():
            _train_ind, _test_ind = [], []
            # Create lpo
            lpo_dict[c] = LeavePAdjacentOut(n, self.stride)
            # Indices set for each class
            y_dict[c] = y_ind_arr[y == c]
            if verbose: print(f"Class {c}: {y_dict[c]}")
        classes.reverse()
        return self._recurse(lpo_dict, y_dict, classes)
    
    def _recurse(self, lpo_dict, y_dict, classes):
        if len(classes) == 1:
            c = classes.pop()
            for train_ind, test_ind in lpo_dict[c].split(y_dict[c]):
                train_ind = y_dict[c][train_ind]
                test_ind = y_dict[c][test_ind]
                yield train_ind, test_ind
        elif len(classes) > 1:
            c = classes.pop()
            to_recurse = list(self._recurse(lpo_dict, y_dict, classes))
            for train_ind, test_ind in lpo_dict[c].split(y_dict[c]):
                train_ind = y_dict[c][train_ind]
                test_ind = y_dict[c][test_ind]
                for res in to_recurse:
                    yield list((*train_ind, *res[0])), list((*test_ind, *res[1]))