1"""Split dataframes into subsets (e.g., train/valid/test)."""
  2
  3# %% IMPORTS
  4
  5import abc
  6import typing as T
  7
  8import numpy as np
  9import numpy.typing as npt
 10import pydantic as pdt
 11from sklearn import model_selection
 12
 13from bikes.core import schemas
 14
 15# %% TYPES
 16
 17Index = npt.NDArray[np.int64]
 18TrainTestIndex = tuple[Index, Index]
 19TrainTestSplits = T.Iterator[TrainTestIndex]
 20
 21# %% SPLITTERS
 22
 23
 24class Splitter(abc.ABC, pdt.BaseModel, strict=True, frozen=True, extra="forbid"):
 25    """Base class for a splitter.
 26
 27    Use splitters to split data in sets.
 28    e.g., split between a train/test subsets.
 29
 30    # https://scikit-learn.org/stable/glossary.html#term-CV-splitter
 31    """
 32
 33    KIND: str
 34
 35    @abc.abstractmethod
 36    def split(
 37        self, inputs: schemas.Inputs, targets: schemas.Targets, groups: Index | None = None
 38    ) -> TrainTestSplits:
 39        """Split a dataframe into subsets.
 40
 41        Args:
 42            inputs (schemas.Inputs): model inputs.
 43            targets (schemas.Targets): model targets.
 44            groups (Index | None, optional): group labels.
 45
 46        Returns:
 47            TrainTestSplits: iterator over the dataframe train/test splits.
 48        """
 49
 50    @abc.abstractmethod
 51    def get_n_splits(
 52        self, inputs: schemas.Inputs, targets: schemas.Targets, groups: Index | None = None
 53    ) -> int:
 54        """Get the number of splits generated.
 55
 56        Args:
 57            inputs (schemas.Inputs): models inputs.
 58            targets (schemas.Targets): model targets.
 59            groups (Index | None, optional): group labels.
 60
 61        Returns:
 62            int: number of splits generated.
 63        """
 64
 65
 66class TrainTestSplitter(Splitter):
 67    """Split a dataframe into a train and test set.
 68
 69    Parameters:
 70        shuffle (bool): shuffle the dataset. Default is False.
 71        test_size (int | float): number/ratio for the test set.
 72        random_state (int): random state for the splitter object.
 73    """
 74
 75    KIND: T.Literal["TrainTestSplitter"] = "TrainTestSplitter"
 76
 77    shuffle: bool = False  # required (time sensitive)
 78    test_size: int | float = 24 * 30 * 2  # 2 months
 79    random_state: int = 42
 80
 81    @T.override
 82    def split(
 83        self, inputs: schemas.Inputs, targets: schemas.Targets, groups: Index | None = None
 84    ) -> TrainTestSplits:
 85        index = np.arange(len(inputs))  # return integer position
 86        train_index, test_index = model_selection.train_test_split(
 87            index, shuffle=self.shuffle, test_size=self.test_size, random_state=self.random_state
 88        )
 89        yield train_index, test_index
 90
 91    @T.override
 92    def get_n_splits(
 93        self, inputs: schemas.Inputs, targets: schemas.Targets, groups: Index | None = None
 94    ) -> int:
 95        return 1
 96
 97
 98class TimeSeriesSplitter(Splitter):
 99    """Split a dataframe into fixed time series subsets.
100
101    Parameters:
102        gap (int): gap between splits.
103        n_splits (int): number of split to generate.
104        test_size (int | float): number or ratio for the test dataset.
105    """
106
107    KIND: T.Literal["TimeSeriesSplitter"] = "TimeSeriesSplitter"
108
109    gap: int = 0
110    n_splits: int = 4
111    test_size: int | float = 24 * 30 * 2  # 2 months
112
113    @T.override
114    def split(
115        self, inputs: schemas.Inputs, targets: schemas.Targets, groups: Index | None = None
116    ) -> TrainTestSplits:
117        splitter = model_selection.TimeSeriesSplit(n_splits=self.n_splits, test_size=self.test_size)
118        yield from splitter.split(inputs)
119
120    @T.override
121    def get_n_splits(
122        self, inputs: schemas.Inputs, targets: schemas.Targets, groups: Index | None = None
123    ) -> int:
124        return self.n_splits
125
126
127SplitterKind = TrainTestSplitter | TimeSeriesSplitter