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,
 38        inputs: schemas.Inputs,
 39        targets: schemas.Targets,
 40        groups: Index | None = None,
 41    ) -> TrainTestSplits:
 42        """Split a dataframe into subsets.
 43
 44        Args:
 45            inputs (schemas.Inputs): model inputs.
 46            targets (schemas.Targets): model targets.
 47            groups (Index | None, optional): group labels.
 48
 49        Returns:
 50            TrainTestSplits: iterator over the dataframe train/test splits.
 51        """
 52
 53    @abc.abstractmethod
 54    def get_n_splits(
 55        self,
 56        inputs: schemas.Inputs,
 57        targets: schemas.Targets,
 58        groups: Index | None = None,
 59    ) -> int:
 60        """Get the number of splits generated.
 61
 62        Args:
 63            inputs (schemas.Inputs): models inputs.
 64            targets (schemas.Targets): model targets.
 65            groups (Index | None, optional): group labels.
 66
 67        Returns:
 68            int: number of splits generated.
 69        """
 70
 71
 72class TrainTestSplitter(Splitter):
 73    """Split a dataframe into a train and test set.
 74
 75    Parameters:
 76        shuffle (bool): shuffle the dataset. Default is False.
 77        test_size (int | float): number/ratio for the test set.
 78        random_state (int): random state for the splitter object.
 79    """
 80
 81    KIND: T.Literal["TrainTestSplitter"] = "TrainTestSplitter"
 82
 83    shuffle: bool = False  # required (time sensitive)
 84    test_size: int | float = 24 * 30 * 2  # 2 months
 85    random_state: int = 42
 86
 87    @T.override
 88    def split(
 89        self,
 90        inputs: schemas.Inputs,
 91        targets: schemas.Targets,
 92        groups: Index | None = None,
 93    ) -> TrainTestSplits:
 94        index = np.arange(len(inputs))  # return integer position
 95        train_index, test_index = model_selection.train_test_split(
 96            index,
 97            shuffle=self.shuffle,
 98            test_size=self.test_size,
 99            random_state=self.random_state,
100        )
101        yield train_index, test_index
102
103    @T.override
104    def get_n_splits(
105        self,
106        inputs: schemas.Inputs,
107        targets: schemas.Targets,
108        groups: Index | None = None,
109    ) -> int:
110        return 1
111
112
113class TimeSeriesSplitter(Splitter):
114    """Split a dataframe into fixed time series subsets.
115
116    Parameters:
117        gap (int): gap between splits.
118        n_splits (int): number of split to generate.
119        test_size (int | float): number or ratio for the test dataset.
120    """
121
122    KIND: T.Literal["TimeSeriesSplitter"] = "TimeSeriesSplitter"
123
124    gap: int = 0
125    n_splits: int = 4
126    test_size: int | float = 24 * 30 * 2  # 2 months
127
128    @T.override
129    def split(
130        self,
131        inputs: schemas.Inputs,
132        targets: schemas.Targets,
133        groups: Index | None = None,
134    ) -> TrainTestSplits:
135        splitter = model_selection.TimeSeriesSplit(
136            n_splits=self.n_splits, test_size=self.test_size, gap=self.gap
137        )
138        yield from splitter.split(inputs)
139
140    @T.override
141    def get_n_splits(
142        self,
143        inputs: schemas.Inputs,
144        targets: schemas.Targets,
145        groups: Index | None = None,
146    ) -> int:
147        return self.n_splits
148
149
150SplitterKind = TrainTestSplitter | TimeSeriesSplitter