{"cells":[{"metadata":{},"cell_type":"markdown","source":"Here is my optimized version of \"[Faster stratified cross-validation (V2)](https://www.kaggle.com/frednavruzov/faster-stratified-cross-validation-v2)\".\n\nIt runs ca. 5 times faster than the original version and has such a plain-dumb logic that it is difficult to make a mistake :) At also has very low memory consumption.\n\nI reused almost the whole original notebook so that you could test yourself.\n\nThe new code is in the function CustomStratifiedKFold2"},{"metadata":{"_uuid":"8f2839f25d086af736a60e9eeb907d3b93b6e0e5","_cell_guid":"b1076dfc-b9ad-4769-8c92-a6c4dae69d19","trusted":true},"cell_type":"code","source":"import pandas as pd\nimport numpy as np\nfrom sklearn.model_selection import StratifiedKFold\nfrom tqdm import tqdm\nfrom copy import deepcopy","execution_count":null,"outputs":[]},{"metadata":{"_uuid":"d629ff2d2480ee46fbb7e2d37f6b5fab8052498a","_cell_guid":"79c7e3d0-c299-4dcb-8224-4455121ee9b0","trusted":true},"cell_type":"code","source":"# let's define our own stratified validator \"with blackjack and hookers\" :)\nclass CustomStratifiedKFold:\n    \"\"\"\n    Faster (yet memory-heavier) stratified cross-validation split\n    Best suited for longer time-series with many different `y` groups\n    \"\"\"\n    def __init__(\n        self,\n        n_splits: int = 5,\n        shuffle: bool = True,\n        random_state: int = 42\n    ):\n        self.n_splits = n_splits\n        self.shuffle = shuffle\n        self.seed = random_state\n        self.folds_ = [(list(), list()) for _ in range(n_splits)]\n        self.randomizer_ = np.random.RandomState(random_state)\n        self.groups_ = None\n        self.counts_ = None\n        self.s_ = None\n\n    def split(self, X, y):\n        sorted_y = pd.Series(y).reset_index(drop=True).sort_values().astype('category').cat.codes\n        self.s_ = pd.Series(data=sorted_y.index.values, index=sorted_y)\n        self.groups_ = self.s_.index.unique()\n        self.counts_ = np.bincount(self.s_.index)\n\n        if self.n_splits > self.counts_.min():\n            raise ValueError(\n                f'Cannot split {self.counts_.min()} elements in smallest group on {self.n_splits} folds'\n            )\n\n        shift = 0\n        for cnt in tqdm(self.counts_, desc='processing unique strats'):\n            # get array of initial data's indices\n            arr = self.s_.iloc[shift:shift + cnt].values\n            # shuffle data if needed\n            if self.shuffle:\n                self.randomizer_.shuffle(arr)\n            folds = np.array_split(arr, self.n_splits)\n            # extend outer folds by elements from micro-folds\n            for i in range(self.n_splits):\n                cp = deepcopy(folds)\n                # extend val indices\n                val_chunk = cp.pop(i).tolist()\n                self.folds_[i][1].extend(val_chunk)\n                # extend train indices\n                if self.shuffle:\n                    cp = self.randomizer_.permutation(cp)\n                train_chunk = np.hstack(cp).tolist()\n                self.folds_[i][0].extend(train_chunk)\n\n            # shift to the next group\n            shift += cnt\n        assert shift == len(self.s_)\n\n        for (t, v) in self.folds_:\n            yield (\n                np.array(self.randomizer_.permutation(t) if self.shuffle else t, dtype=np.int32),\n                np.array(self.randomizer_.permutation(v) if self.shuffle else v, dtype=np.int32)\n            )","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"from collections import defaultdict\n\nclass CustomStratifiedKFold2:\n    def __init__(self, n_splits: int = 5, shuffle: bool = True, random_state: int = 42):\n        self.n_splits    = n_splits\n        self.shuffle     = shuffle\n        self.randomizer_ = np.random.RandomState(random_state)\n\n    def split(self, X, y):\n        folds = defaultdict(lambda: np.zeros(len(y), dtype=np.int8))              # global fold reference\n\n        for i, (g, u) in enumerate(tqdm(pd.Series(y).groupby(by=y), desc='processing strats')):\n            strata = u.index.values                                               # array of initial data's indices for the given strata\n\n            l = len(strata)\n            if self.n_splits > l:\n                raise ValueError(f'Cannot split {l} elements in group {u} on {self.n_splits} folds')\n            \n            if self.shuffle:                                                      # shuffle if asked so\n                self.randomizer_.shuffle(strata)\n                \n            tl = l // self.n_splits\n            for j in range(self.n_splits):\n                folds[j][strata]                =  1                              # first set the whole strata to train\n                folds[j][strata[tl*j:tl*(j+1)]] =  2                              # now set the valid part \n                \n        for i in range(self.n_splits):\n            yield (folds[i] == 1, folds[i] == 2)","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"### Let's create some fake data to test different approaches on"},{"metadata":{"trusted":true},"cell_type":"code","source":"N_FOLDS = 5\nSEED = 42\nSHUFFLE = True\nNUM_UNIQUES = 1000\nN = 7000000\n\nrandomizer = np.random.RandomState(SEED)\nstrat_column = pd.Series(randomizer.randint(0, NUM_UNIQUES, N, dtype=np.int32))\nprint(strat_column.nunique())\nstrat_column.head(10)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"%%time\n# let's check usual StratifiedKFold speed\nskf = StratifiedKFold(n_splits=N_FOLDS, shuffle=SHUFFLE, random_state=SEED)\n\nfolds = list(skf.split(\n    # we don't actually need `X` to produce indices, only `y`\n    X=np.zeros(len(strat_column)),\n    y=strat_column\n))","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"# let's check whether class balance is preserved\nprint('train')\nprint(strat_column.iloc[folds[0][0]].value_counts(normalize=True).sort_index())\nprint('val')\nprint(strat_column.iloc[folds[0][1]].value_counts(normalize=True).sort_index())","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"# check all indices are there in joined validation blocks\nassert len(set(np.hstack([v for (tr,v) in folds]).tolist())) == len(strat_column)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"%%time\n# let's check updated StratifiedKFold speed\nskf = CustomStratifiedKFold(n_splits=N_FOLDS, shuffle=SHUFFLE, random_state=SEED)\n\nfolds = list(skf.split(\n    # we don't actually need `X` to produce indices, only `y`\n    X=None,\n    y=strat_column\n))","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":" Well, even on smaller dataset we get** 3x+** speed improvement"},{"metadata":{"trusted":true},"cell_type":"code","source":"# let's check whether class balance is preserved also in new method\nprint('train')\nprint(strat_column.iloc[folds[1][0]].value_counts(normalize=True).sort_index())\nprint('val')\nprint(strat_column.iloc[folds[1][1]].value_counts(normalize=True).sort_index())","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"# check all indices are there in joined validation blocks (also for the new strategy)\nassert len(set(np.hstack([v for (tr,v) in folds]).tolist())) == len(strat_column)","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"### let's create real-world example"},{"metadata":{"trusted":true},"cell_type":"code","source":"N_FOLDS = 5\nSEED = 42\nSHUFFLE = True\nNUM_UNIQUES = 10000\nN = 20000000\n\nstrat_column = pd.Series(randomizer.randint(0, NUM_UNIQUES, N, dtype=np.int32))\nprint(strat_column.nunique())\nstrat_column.head(10)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"%%time\n# let's check updated StratifiedKFold speed on heavier task (however, notice rapid memory peak)\n# don't try to run this with usual `StratifiedKFold` or prepare to wait A LOT\nskf = CustomStratifiedKFold(n_splits=N_FOLDS, shuffle=SHUFFLE, random_state=SEED)\n\nfolds = list(skf.split(\n    # we don't actually need `X` to produce indices, only `y`\n    X=None,\n    y=strat_column\n))","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"# let's check whether class balance is preserved also in new method\nprint('train')\nprint(strat_column.iloc[folds[1][0]].value_counts(normalize=True).sort_index())\nprint('val')\nprint(strat_column.iloc[folds[1][1]].value_counts(normalize=True).sort_index())","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"Well, almost as fast (or even faster at commiting mode) as `StratifiedKFold` on much smaller/less diverse dataset!"},{"metadata":{"trusted":true},"cell_type":"code","source":"%%time\nskf = CustomStratifiedKFold2(n_splits=N_FOLDS, shuffle=SHUFFLE, random_state=SEED)\n\nfolds = list(skf.split(\n    # we don't actually need `X` to produce indices, only `y`\n    X=None,\n    y=strat_column\n))","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"# let's check whether class balance is preserved also in new method\nprint('train')\nprint(strat_column.iloc[folds[1][0]].value_counts(normalize=True).sort_index())\nprint('val')\nprint(strat_column.iloc[folds[1][1]].value_counts(normalize=True).sort_index())","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"","execution_count":null,"outputs":[]}],"metadata":{"kernelspec":{"language":"python","display_name":"Python 3","name":"python3"},"language_info":{"pygments_lexer":"ipython3","nbconvert_exporter":"python","version":"3.6.4","file_extension":".py","codemirror_mode":{"name":"ipython","version":3},"name":"python","mimetype":"text/x-python"}},"nbformat":4,"nbformat_minor":1}