# AUTOGENERATED! DO NOT EDIT! File to edit: ../../nbs/013_data.core.ipynb.

# %% auto 0
__all__ = ['TSCategorize', 'TSRegression', 'TSForecasting', 'get_tsimage_dls', 'NumpyTensor', 'ToNumpyTensor', 'TSTensor',
           'ToTSTensor', 'show_tuple', 'TSLabelTensor', 'TSMaskTensor', 'ToFloat', 'ToInt', 'TSClassification',
           'TSMultiLabelClassification', 'NumpyTensorBlock', 'TSTensorBlock', 'TorchDataset', 'NumpyDataset',
           'TSDataset', 'NoTfmLists', 'TSTfmdLists', 'NumpyDatasets', 'tscoll_repr', 'TSDatasets', 'add_ds',
           'NumpyDataLoader', 'TSDataLoader', 'NumpyDataLoaders', 'TSDataLoaders', 'StratifiedSampler', 'get_c',
           'get_best_dl_params', 'get_best_dls_params', 'get_ts_dls', 'get_ts_dl', 'get_subset_dl',
           'get_time_per_batch', 'get_dl_percent_per_epoch']

# %% ../../nbs/013_data.core.ipynb 3
from ..imports import *
from sklearn.model_selection import StratifiedKFold
from types import MethodType
from matplotlib.ticker import PercentFormatter
import matplotlib.colors as mcolors
from matplotlib import rcParams
from fastcore.transform import Transform, DisplayedTransform, Pipeline
from fastai.data.transforms import (RandomSplitter, ItemGetter, MultiCategorize, CategoryMap,
Category, MultiCategory, Categorize)
from fastai.data.core import TfmdLists, Datasets, TfmdDL, DataLoaders
from fastai.data.load import DataLoader
from fastai.data.block import CategoryBlock, DataBlock
from fastai.losses import MSELossFlat, CrossEntropyLossFlat, BCEWithLogitsLossFlat
from fastai.vision.data import get_grid
from ..utils import *

# %% ../../nbs/013_data.core.ipynb 6
class NumpyTensor(TensorBase):
    "Returns a `tensor` with subclass `NumpyTensor` that has a show method"

    def __new__(cls, o, dtype=None, device=None, **kwargs):
        if dtype is not None or device is not None: o = torch.as_tensor(o, dtype=dtype, device=device)
        res = cast(o, cls)
        for k,v in kwargs.items(): setattr(res, k, v)
        return res
        
    @property
    def data(self): return cast(self, Tensor)
    
    def __repr__(self):
        if self.ndim > 0: return f'NumpyTensor(shape:{tuple(self.shape)}, device={self.device}, dtype={self.dtype})'
        else: return f'NumpyTensor([{self}], device={self.device}, dtype={self.dtype})'
        

    def show(self, ax=None, ctx=None, title=None, **kwargs):
        if self.ndim == 0: return str(self)
        elif self.ndim != 2: self = type(self)(to2d(self))
        if not isinstance(self,np.ndarray): self = self.detach().cpu().numpy()
        ax = ifnone(ax, ctx)
        if ax is None: _, ax = plt.subplots(**kwargs)
        ax.plot(self.T)
        ax.axis(xmin=0, xmax=self.shape[-1] - 1)
        title_color = kwargs['title_color'] if 'title_color' in kwargs else rcParams['axes.labelcolor']
        if title is not None:
            if is_listy(title): title = str(title)[1:-1]
            ax.set_title(title, weight='bold', color=title_color)
        plt.tight_layout()
        return ax
    

class ToNumpyTensor(Transform):
    "Transforms an object into NumpyTensor"
    def encodes(self, o): return NumpyTensor(o)

# %% ../../nbs/013_data.core.ipynb 7
class TSTensor(NumpyTensor):
    '''Returns a `tensor` with subclass `TSTensor` that has a show method'''

    @property
    def vars(self):
        return self.shape[-2]

    @property
    def len(self): return self.shape[-1]

    def __repr__(self):
        if self.ndim > 3:
            return f'TSTensor(shape:{tuple(self.shape)}, device={self.device}, dtype={self.dtype})'
        elif self.ndim == 3:
            return f'TSTensor(samples:{self.shape[-3]}, vars:{self.shape[-2]}, len:{self.shape[-1]}, device={self.device}, dtype={self.dtype})'
        elif self.ndim == 2:
            return f'TSTensor(vars:{self.shape[-2]}, len:{self.shape[-1]}, device={self.device}, dtype={self.dtype})'
        elif self.ndim == 1:
            return f'TSTensor(len:{self.shape[-1]}, device={self.device}, dtype={self.dtype})'
        else: return f'TSTensor([{self}], device={self.device}, dtype={self.dtype})'


class ToTSTensor(Transform):
    "Transforms an object into TSTensor"
    def encodes(self, o): return TSTensor(o)
    
    
@delegates(plt.subplots)
def show_tuple(tup, **kwargs):
    "Display a timeseries plot from a decoded tuple"
    if len(tup) == 1: title = 'unlabeled'
    elif is_listy(tup[1]): title = str(tup[1])[1:-1]
    else: title = str(tup[1])
    tup[0].show(title=title, **kwargs)

# %% ../../nbs/013_data.core.ipynb 21
class TSLabelTensor(NumpyTensor): 
    def __repr__(self):
        if self.ndim == 0: return f'{self}'
        else: return f'TSLabelTensor(shape:{tuple(self.shape)}, device={self.device}, dtype={self.dtype})'

class TSMaskTensor(NumpyTensor): 
    def __repr__(self):
        if self.ndim == 0: return f'{self}'
        else: return f'TSMaskTensor(shape:{tuple(self.shape)}, device={self.device}, dtype={self.dtype})'

# %% ../../nbs/013_data.core.ipynb 24
class ToFloat(Transform):
    "Transforms an object dtype to float (vectorized)"
    vectorized=True
    loss_func=MSELossFlat()
    def encodes(self, o:torch.Tensor): return o.float()
    def encodes(self, o): return o.astype(np.float32)
    def decodes(self, o): 
        if o.ndim==0: return TitledFloat(o) 
        else: 
            return TitledTuple(o.cpu().numpy().tolist())
            

class ToInt(Transform):
    "Transforms an object dtype to int"
    def encodes(self, o:torch.Tensor): return o.long()
    def encodes(self, o): return o.astype(np.float32).astype(np.int64)
    def decodes(self, o): 
        if o.ndim==0: return TitledFloat(o) 
        else: 
            return TitledTuple(o.cpu().numpy().tolist())
    
    
class TSClassification(DisplayedTransform):
    "Vectorized, reversible transform of category string to `vocab` id"
    loss_func,order,vectorized=CrossEntropyLossFlat(),1,True

    def __init__(self, vocab=None, sort=True):
        if vocab is not None: vocab = CategoryMap(vocab, sort=sort, add_na=False)
        store_attr()

    def setups(self, dset):
        if self.vocab is None and dset is not None:
            dset = np.asarray(dset).flatten()
            self.vocab = CategoryMap(dset, sort=self.sort, add_na=False)
        self.c = len(self.vocab)
        self.vocab_keys = np.array(list(self.vocab.o2i.keys()))[:, None]

    def encodes(self, o):
        try:
            if not hasattr(o, "ndim") or o.ndim <= 1:
                return TensorCategory((self.vocab_keys == o).argmax(axis=0))
            else:
                return TensorCategory((self.vocab_keys == o.flatten()).argmax(axis=0).reshape(*o.shape))
        except KeyError as e:
            raise KeyError(f"Label '{o}' was not included in the training dataset") from e
    def decodes(self, o):
        if not is_iter(o):
            return Category(self.vocab[o])
        elif o.ndim <= 1:
            return stack([Category(self.vocab[oi]) for oi in o])
        else:
            return stack(MultiCategory(self.vocab[o.flatten()])).reshape(*o.shape)

        
TSCategorize = TSClassification
TSRegression = ToFloat
TSForecasting = ToFloat

# %% ../../nbs/013_data.core.ipynb 29
class TSMultiLabelClassification(Categorize):
    "Reversible combined transform of multi-category strings to one-hot encoded `vocab` id"
    loss_func,order=BCEWithLogitsLossFlat(),1
    def __init__(self, c=None, vocab=None, add_na=False, sort=True): 
        super().__init__(vocab=vocab,add_na=add_na,sort=sort)
        self.c = c

    def setups(self, dsets):
        if not dsets: return
        if self.vocab is None:
            vals = set()
            for b in dsets: vals = vals.union(set(b))
            self.vocab = CategoryMap(list(vals), add_na=self.add_na)
            if self.c is None: self.c = len(self.vocab)
            if not self.c: warn("Couldn't infer the number of classes, please pass a value for `c` at init")

    def encodes(self, o):
        if not all(elem in self.vocab.o2i.keys() for elem in o):
            diff = [elem for elem in o if elem not in self.vocab.o2i.keys()]
            diff_str = "', '".join(diff)
            raise KeyError(f"Labels '{diff_str}' were not included in the training dataset")
        return TensorMultiCategory(one_hot([self.vocab.o2i[o_] for o_ in o], self.c).float())
    def decodes(self, o): return MultiCategory([self.vocab[o_] for o_ in one_hot_decode(o, None)])

# %% ../../nbs/013_data.core.ipynb 30
class NumpyTensorBlock():
    def __init__(self, type_tfms=None, item_tfms=None, batch_tfms=None, dl_type=None, dls_kwargs=None):
        self.type_tfms  =                 L(type_tfms)
        self.item_tfms  = ToNumpyTensor + L(item_tfms)
        self.batch_tfms =                 L(batch_tfms)
        self.dl_type,self.dls_kwargs = dl_type,({} if dls_kwargs is None else dls_kwargs)
        
class TSTensorBlock():
    def __init__(self, type_tfms=None, item_tfms=None, batch_tfms=None, dl_type=None, dls_kwargs=None):
        self.type_tfms  =              L(type_tfms)
        self.item_tfms  = ToTSTensor + L(item_tfms)
        self.batch_tfms =              L(batch_tfms)
        self.dl_type,self.dls_kwargs = dl_type,({} if dls_kwargs is None else dls_kwargs)

# %% ../../nbs/013_data.core.ipynb 32
class TorchDataset():
    def __init__(self, X, y=None): self.X, self.y = X, y
    def __getitem__(self, idx): return (self.X[idx],) if self.y is None else (self.X[idx], self.y[idx])
    def __len__(self): return len(self.X)

    
class NumpyDataset():
    def __init__(self, X, y=None, types=None): self.X, self.y, self.types = X, y, types
    def __getitem__(self, idx):
        if self.types is None: return (self.X[idx], self.y[idx]) if self.y is not None else (self.X[idx])
        else: return (self.types[0](self.X[idx]), self.types[1](self.y[idx])) if self.y is not None else (self.types[0](self.X[idx]))
    def __len__(self): return len(self.X)


class TSDataset():
    _types = TSTensor, TSLabelTensor
    def __init__(self, X, y=None, split=None, sel_vars=None, sel_steps=None, types=None, dtype=None, device=None):
        self.X, self.y, self.split = X, y, split
        self.sel_vars = sel_vars
        self.sel_steps = sel_steps
        self.multi_idx = sel_vars is not None or sel_steps is not None
        if types is not None: self._types = listify(types)
        self.dtype, self.device = dtype, device
    def __getitem__(self, idx):
        if self.split is not None:
            idx = self.split[idx]
        if hasattr(self.X, 'oindex'):
            X = self._types[0](self.X.oindex[idx, self.sel_vars, self.sel_steps] if self.multi_idx \
                               else self.X.oindex[idx], device=self.device, dtype=self.dtype)
        elif hasattr(self.X, 'compute'):
            X = self._types[0](self.X[idx, self.sel_vars, self.sel_steps].compute() if self.multi_idx \
                               else self.X[idx].compute(), device=self.device, dtype=self.dtype)
        else:
            X = self._types[0](self.X[idx, self.sel_vars, self.sel_steps] if self.multi_idx \
                               else self.X[idx], device=self.device, dtype=self.dtype)
        if self.y is None:
            return (X, )
        if hasattr(self.y, 'oindex'):
            y = self._types[1](self.y.oindex[idx], device=self.device, dtype=self.dtype)
        elif hasattr(self.X, 'compute'):
            y = self._types[1](self.y[idx].compute(), device=self.device, dtype=self.dtype)
        else:
            y = self._types[1](self.y[idx], device=self.device, dtype=self.dtype)
        return (X, y)
    def __len__(self): return len(self.X) if self.split is None else len(self.split)

# %% ../../nbs/013_data.core.ipynb 34
def _flatten_list(l):
    if not is_listy(l) or len(l) == 0: return l
    return [item for sublist in l for item in listify(sublist)]

def _remove_brackets(l):
    return [li if (not li or not is_listy(li) or len(li) > 1) else li[0] for li in l]
    
class NoTfmLists(TfmdLists):
    def __init__(self, items, tfms=None, splits=None, split_idx=None, types=None, do_setup=False, **kwargs):
        self.splits = ifnone(splits, L(np.arange(len(items)).tolist(),[]))
        self._splits = np.asarray(_flatten_list(self.splits))
        store_attr('items,types,split_idx')
        self.tfms = Pipeline(split_idx=split_idx)
    def subset(self, i, **kwargs): return type(self)(self.items, splits=self.splits[i], split_idx=i, do_setup=False, types=self.types, 
                                                     **kwargs)
    def __getitem__(self, it):
        if hasattr(self.items, 'oindex'): return self.items.oindex[self._splits[it]]
        else: return self.items[self._splits[it]]
    def __len__(self): return len(self._splits)
    def __repr__(self): 
        if hasattr(self.items, "shape"):
            return f"{self.__class__.__name__}: {self.items.__class__.__name__}{(len(self), *self.items.shape[1:])}"
        else:
            return f"{self.__class__.__name__}: {self.items.__class__.__name__}(0)"
    def _new(self, items, split_idx=None, **kwargs):
        split_idx = ifnone(split_idx, self.split_idx)
        return type(self)(items, split_idx=split_idx, do_setup=False, types=self.types, **kwargs)
    def decode(self, o, **kwargs): return o
    def new_empty(self): return self._new([])

NoTfmLists.train, NoTfmLists.valid = add_props(lambda i,x: x.subset(i))

class TSTfmdLists(TfmdLists):
    def __getitem__(self, it):
        # res = self._get(it)
        if hasattr(self.items, 'oindex'): res = self.items.oindex[it] 
        else: res = self.items[it]
        if self._after_item is None: return res
        else: return self._after_item(res)

# %% ../../nbs/013_data.core.ipynb 41
@delegates(Datasets.__init__)
class NumpyDatasets(Datasets):
    "A dataset that creates tuples from X (and y) and applies `tfms` of type item_tfms"
    typs = NumpyTensor,tensor
    def __init__(self, X=None, y=None, items=None, tfms=None, tls=None, n_inp=None, dl_type=None, inplace=True, **kwargs):

        self.tfms, self.inplace = tfms, inplace

        if X is not None:
            if not hasattr(X, '__array__'): X = np.asarray(X)
            elif hasattr(X, "iloc"): X = to3d(X)
        if y is not None:
            if not hasattr(y, '__array__'):  y = np.asarray(y)
            elif hasattr(y, "iloc"): y = toarray(y)

        if tls is None:
            items = tuple((X,)) if y is None else tuple((X, y))

            if tfms is None:
                self.tfms, lts = [None] * len(items), [NoTfmLists] * len(items)
            else:
                self.tfms = _remove_brackets(tfms)
                ts = [NoTfmLists if t is None else TSTfmdLists if getattr(t, 'vectorized', None) else TfmdLists for t in self.tfms]

            self.tls = L(lt(item, t, **kwargs) for lt,item,t in zip(lts, items, self.tfms))
            if len(self.tls) > 0 and len(self.tls[0]) > 0:
                self.typs = [type(tl[0]) if isinstance(tl[0], torch.Tensor) else self.typs[i] for i,tl in enumerate(self.tls)]
            self.ptls = L([typ(stack(tl[:])) for i,(tl,typ) in enumerate(zip(self.tls,self.typs))]) if inplace else self.tls
        else:
            self.tls = tls
            if len(self.tls) > 0 and len(self.tls[0]) > 0:
                self.typs = [type(tl[0]) if isinstance(tl[0], torch.Tensor) else self.typs[i] for i,tl in enumerate(self.tls)]
            self.ptls = L([typ(stack(tl[:])) for i,(tl,typ) in enumerate(zip(self.tls,self.typs))]) if inplace and len(tls[0]) != 0 else tls

        self.n_inp = 1
        if 'splits' in kwargs:
            split_idxs = kwargs['splits']
            try: split_idxs = flatten_list(split_idxs)
            except: pass
        else: split_idxs = L(np.arange(len(self.tls[0])).tolist())
        self.split_idxs = L(split_idxs)

    def __getitem__(self, it):
        if self.inplace:
            return tuple([ptl[it] for ptl in self.ptls])
        else:
            return tuple([typ(stack(ptl[it])) for i,(ptl,typ) in enumerate(zip(self.ptls,self.typs))])

    def subset(self, i):
        if is_indexer(i):
            return type(self)(tls=L([tl.subset(i) for tl in self.tls]), inplace=self.inplace, tfms=self.tfms,
                              splits=None if self.splits is None else self.splits[i], split_idx=i)
        else:
            splits = None if self.splits is None else L(np.arange(len(i)).tolist())
            return type(self)(*self[i], inplace=True, tfms=None, splits=splits, split_idx=ifnone(self.split_idx, 1))

    def __len__(self): return len(self.tls[0])
    
    def _new(self, X, y=None, **kwargs):
        return type(self)(X, y=y, tfms=self.tfms, inplace=self.inplace, do_setup=False, **kwargs)

    def new_empty(self): return type(self)(tls=[tl.new_empty() for tl in self.tls], n_inp=self.n_inp, inplace=self.inplace)
    
    def show_at(self, idx, **kwargs):
        self.show(self[idx], **kwargs)
        plt.show()
        
    def __repr__(self): return tscoll_repr(self)


def tscoll_repr(c, max_n=10):
    "String repr of up to `max_n` items of (possibly lazy) collection `c`"
    _len = len(c)
    if _len == 0: return coll_repr(c)
    return f'(#{_len}) {L(c[i] for i in range(min(len(c), max_n)))} ...]'

# %% ../../nbs/013_data.core.ipynb 42
@delegates(NumpyDatasets.__init__)
class TSDatasets(NumpyDatasets):
    """A dataset that creates tuples from X (and optionally y) and applies `item_tfms`"""
    typs = TSTensor,tensor
    def __init__(self, X=None, y=None, items=None, sel_vars=None, sel_steps=None, tfms=None, tls=None, n_inp=None, dl_type=None, 
                 inplace=True, **kwargs):

        self.multi_index = False
        if sel_vars is None or (type(sel_vars) == slice and sel_vars == slice(None)): self.sel_vars = slice(None)
        elif type(sel_vars) == slice: 
            self.sel_vars = sel_vars
            self.multi_index = True
        else:
            self.sel_vars = np.asarray(sel_vars)
            if sel_steps is not None and type(sel_steps) != slice: self.sel_vars = sel_vars[:, None]
            self.multi_index = True
        if sel_steps is None or (type(sel_steps) == slice and sel_steps == slice(None)): self.sel_steps = slice(None)
        elif type(sel_steps) == slice: 
            self.sel_steps = sel_steps
            self.multi_index = True
        else: 
            self.sel_steps = np.asarray(sel_steps)
            self.multi_index = True
        self.tfms, self.inplace = tfms, inplace

        if X is not None:
            if not hasattr(X, '__array__'): X = np.asarray(X)
            X = to3d(X)
        if y is not None:
            if not hasattr(y, '__array__'):  y = np.asarray(y)
            elif hasattr(y, "iloc"): y = toarray(y)

        if tls is None:
            items = tuple((X,)) if y is None else tuple((X, y))

            if tfms is None:
                self.tfms, lts = [None] * len(items), [NoTfmLists] * len(items)
            else:
                self.tfms = _remove_brackets(tfms)
                lts = [NoTfmLists if t is None else TSTfmdLists if getattr(t, 'vectorized', None) else TfmdLists for t in self.tfms]

            self.tls = L(lt(item, t, **kwargs) for lt,item,t in zip(lts, items, self.tfms))
            if len(self.tls) > 0 and len(self.tls[0]) > 0:
                self.typs = [type(tl[0]) if isinstance(tl[0], torch.Tensor) else self.typs[i] for i,tl in enumerate(self.tls)]
            self.ptls = L([typ(stack(tl[:]))[...,self.sel_vars, self.sel_steps] if (i==0 and self.multi_index) else typ(stack(tl[:])) \
                            for i,(tl,typ) in enumerate(zip(self.tls,self.typs))]) if inplace else self.tls
        else:
            self.tls = tls
            if len(self.tls) > 0 and len(self.tls[0]) > 0:
                self.typs = [type(tl[0]) if isinstance(tl[0], torch.Tensor) else self.typs[i] for i,tl in enumerate(self.tls)]
            self.ptls = L([typ(stack(tl[:]))[...,self.sel_vars, self.sel_steps] if (i==0 and self.multi_index) else typ(stack(tl[:])) \
                            for i,(tl,typ) in enumerate(zip(self.tls,self.typs))]) if inplace and len(tls[0]) != 0 else tls

        self.n_inp = 1
        if 'splits' in kwargs:
            split_idxs = kwargs['splits']
            try: split_idxs = flatten_list(split_idxs)
            except: pass
        else: split_idxs = L(np.arange(len(self.tls[0])).tolist())
        self.split_idxs = L(split_idxs)

    def __getitem__(self, it):
        if self.inplace:
            return tuple([ptl[it] for ptl in self.ptls])
        else:
            return tuple([typ(stack(ptl[it]))[...,self.sel_vars, self.sel_steps] if (i==0 and self.multi_index) else typ(stack(ptl[it])) \
                          for i,(ptl,typ) in enumerate(zip(self.ptls,self.typs))])

    def subset(self, i):
        if is_indexer(i):
            return type(self)(tls=L([tl.subset(i) for tl in self.tls]), inplace=self.inplace, tfms=self.tfms,
                              sel_vars=self.sel_vars, sel_steps=self.sel_steps, splits=None if self.splits is None else self.splits[i], 
                              split_idx=i)
        else:
            splits = None if self.splits is None else L(np.arange(len(i)).tolist())
            return type(self)(*self[i], inplace=True, tfms=None, 
                              sel_vars=self.sel_vars, sel_steps=self.sel_steps, splits=splits, split_idx=ifnone(self.split_idx, 1))
    
    def _new(self, X, y=None, **kwargs):
        return type(self)(X, y=y, sel_vars=self.sel_vars, sel_steps=self.sel_steps, tfms=self.tfms, inplace=self.inplace, 
                          do_setup=False, **kwargs)
    
    def new_empty(self): return type(self)(tls=[tl.new_empty() for tl in self.tls], sel_vars=self.sel_vars, sel_steps=self.sel_steps, 
                                           n_inp=self.n_inp, inplace=self.inplace)

# %% ../../nbs/013_data.core.ipynb 44
def add_ds(dsets, X, y=None, inplace=True):
    "Create test datasets from X (and y) using validation transforms of `dsets`"
    items = tuple((X,)) if y is None else tuple((X, y))
    with_labels = False if y is None else True 
    if isinstance(dsets, TSDatasets):
        tls = dsets.tls if with_labels else dsets.tls[:dsets.n_inp]
        new_tls = L([tl._new(item, split_idx=1) for tl,item in zip(tls, items)])
        return type(dsets)(tls=new_tls, sel_vars=dsets.sel_vars, sel_steps=dsets.sel_steps, inplace=dsets.inplace)
    elif isinstance(dsets, NumpyDatasets):
        tls = dsets.tls if with_labels else dsets.tls[:dsets.n_inp]
        new_tls = L([tl._new(item, split_idx=1) for tl,item in zip(tls, items)])
        return type(dsets)(tls=new_tls, inplace=dsets.inplace)
    elif isinstance(dsets, Datasets):
        tls = dsets.tls if with_labels else dsets.tls[:dsets.n_inp]
        new_tls = L([tl._new(item, split_idx=1) for tl,item in zip(tls, items)])
        return type(dsets)(tls=new_tls)
    elif isinstance(dsets, TfmdLists):
        new_tl = dsets._new(items, split_idx=1)
        return new_tl
    else: 
        raise Exception(f"Expected a `Datasets` or a `TfmdLists` but got {dsets.__class__.__name__}")

@patch
def add_dataset(self:NumpyDatasets, X, y=None, inplace=True):
    return add_ds(self, X, y=y, inplace=inplace)

@patch
def add_test(self:NumpyDatasets, X, y=None, inplace=True):
    return add_ds(self, X, y=y, inplace=inplace)

@patch
def add_unlabeled(self:NumpyDatasets, X, inplace=True):
    return add_ds(self, X, y=None, inplace=inplace)

# %% ../../nbs/013_data.core.ipynb 60
@patch
def _one_pass(self:TfmdDL):
    b = self.do_batch([self.do_item(0)])
    if self.device is not None: b = to_device(b, self.device)
    its = self.after_batch(b)
    self._n_inp = 1 if not isinstance(its, (list,tuple)) or len(its)==1 else len(its)-1
    self._types = explode_types(its)

# %% ../../nbs/013_data.core.ipynb 61
_batch_tfms = ('after_item','before_batch','after_batch')

@delegates(TfmdDL.__init__)
class NumpyDataLoader(TfmdDL):
    idxs = None
    do_item = noops # create batch returns indices
    def __init__(self, dataset, bs=64, shuffle=False, drop_last=False, num_workers=0, verbose=False, do_setup=True, 
                 sort=False, weights=None, partial_n=None, sampler=None, **kwargs):

        if num_workers is None: num_workers = min(16, defaults.cpus)
        if sampler is not None and shuffle:
            raise ValueError('sampler option is mutually exclusive with shuffle')
        
        for nm in _batch_tfms:
            if nm == 'after_batch' and kwargs.get('batch_tfms',None) is not None: kwargs[nm] = Pipeline(listify(kwargs.get('batch_tfms')))
            else: kwargs[nm] = Pipeline(kwargs.get(nm,None))
        bs = min(bs, len(dataset))
        if is_listy(partial_n): partial_n = partial_n[0]
        if isinstance(partial_n, float): partial_n = int(round(partial_n * len(dataset)))
        if partial_n is not None: bs = min(bs, partial_n)
        if weights is not None: weights = weights / weights.sum()
        self.weights, self.partial_n, self.sampler, self.sort, self.do_setup = weights, partial_n, sampler, sort, do_setup
        super().__init__(dataset, bs=bs, shuffle=shuffle, drop_last=drop_last, num_workers=num_workers, verbose=verbose, do_setup=do_setup, **kwargs)
    
    def new_dl(self, X, y=None, bs=64):
        assert X.ndim == 3, "You must pass an X with 3 dimensions [batch_size x n_vars x seq_len]"
        if y is not None: 
            y = listify(y)
        ds = self.dataset.add_dataset(X, y=y)
        return self.new(ds, bs=min(bs, len(X)))

    def create_batch(self, b):
        if self.shuffle or self.sampler is not None:
            if self.sort and hasattr(b, 'sort'): b.sort()
            self.idxs = L(b)
        else: 
            if self.n is not None:
                b = slice(b[0], min(self.n, b[0] + self.bs))
            else:
                b = slice(b[0], b[0] + self.bs)
                
            self.idxs = b
        if hasattr(self, "split_idxs"): self.input_idxs = self.split_idxs[b]
        else: self.input_idxs = self.idxs
        return self.dataset[b]

    def create_item(self, s):
        if self.indexed: return self.dataset[s or 0]
        elif s is None:  return next(self.it)
        else: raise IndexError("Cannot index an iterable dataset numerically - must use `None`.")
    
    def get_idxs(self):
        if self.n==0: return []
        if self.partial_n is not None: n = min(self.partial_n, self.n)
        else: n = self.n
        if self.sampler is not None:
            idxs = np.array(list(iter(self.sampler)))
            if self.partial_n is not None: idxs = idxs[:n]
            return idxs
        if self.weights is not None:
            return np.random.choice(self.n, n, p=self.weights)
        idxs = Inf.count if self.indexed else Inf.nones
        if self.n is not None:
            idxs = np.arange(self.n)
            if self.partial_n is not None:
                idxs = np.random.choice(idxs, n, False)
        if self.shuffle: idxs = self.shuffle_fn(idxs)
        return idxs


    def shuffle_fn(self, idxs):
        return np.random.permutation(idxs)

    def unique_batch(self, max_n=9):
        old_bs = self.bs
        self.bs = 1
        old_get_idxs = self.get_idxs
        self.get_idxs = lambda: Inf.zeros
        out_len = len(self.items)
        types = self.dataset.types
        x, y = [], []
        for _ in range(max_n):
            out = self.one_batch()
            if out_len == 2:
                x.extend(out[0])
                y.extend(out[1])
            else:
                x.extend(out)
        b = (types[0](stack(x)), types[1](stack(y))) if out_len == 2 else (types[0](stack(x)), )
        self.bs = old_bs
        self.get_idxs = old_get_idxs
        return b

    def __len__(self):
        if self.partial_n is None: return super().__len__()
        return self.partial_n//self.bs + (0 if self.drop_last or self.partial_n%self.bs==0 else 1)
    
    @delegates(plt.subplots)
    def show_batch(self, b=None, ctxs=None, max_n=9, nrows=3, ncols=3, figsize=None, unique=False, sharex=True, sharey=False, decode=False, 
                   show_title=True, **kwargs):
        
        old_sort = self.sort 
        self.sort = False # disable sorting when showing a batch to ensure varied samples

        if unique:
            b = self.unique_batch(max_n=max_n)
            sharex, sharey = True, True
        elif b is None: b = self.one_batch()
        if not decode:                                        # decode = False allows you to see the data as seen by the model
            after_batch = self.after_batch
            self.after_batch = Pipeline()
            db = self.decode_batch(b, max_n=max_n)
            self.after_batch = after_batch
        else:
            db = self.decode_batch(b, max_n=max_n)
        ncols = min(ncols, math.ceil(len(db) / ncols))
        nrows = min(nrows, math.ceil(len(db) / ncols))
        max_n = min(max_n, len(db), nrows*ncols)
        if figsize is None: figsize = (ncols*6, math.ceil(max_n/ncols)*4)
        if ctxs is None: ctxs = get_grid(max_n, nrows=nrows, ncols=ncols, figsize=figsize, sharex=sharex, sharey=sharey, **kwargs)
        if show_title:
            for i,ctx in enumerate(ctxs): 
                show_tuple(db[i], ctx=ctx)
        else:
            db = [x for x,_ in db]
            for i,ctx in enumerate(ctxs): 
                db[i].show(ctx=ctx)
        
        self.sort = old_sort

    @delegates(plt.subplots)
    def show_results(self, b, preds, ctxs=None, max_n=9, nrows=3, ncols=3, figsize=None, **kwargs):
        t = self.decode_batch(b, max_n=max_n)
        p = self.decode_batch((b[0],preds), max_n=max_n)
        if figsize is None: figsize = (ncols*6, max_n//ncols*4)
        if ctxs is None: ctxs = get_grid(min(len(t), nrows*ncols), nrows=None, ncols=ncols, figsize=figsize, **kwargs)
        for i,ctx in enumerate(ctxs):
            title = f'True: {t[i][1]}\nPred: {p[i][1]}'
            color = 'green' if t[i][1] == p[i][1] else 'red'
            t[i][0].show(ctx=ctx, title=title, title_color=color)

    @delegates(plt.subplots)
    def show_dist(self, figsize=None, color=None, **kwargs):
        if self.c == 0:
            print('\nunlabeled dataset.\n')
            return
        b = self.one_batch()
        i = getattr(self, 'n_inp', 1 if len(b)==1 else len(b)-1)
        yb = b[i:][0]
        if color == "random": color = random_shuffle(L(mcolors.CSS4_COLORS.keys()))
        elif color is None: color = ['m', 'orange', 'darkblue', 'lightgray']
        figsize = ifnone(figsize, (8, 6))
        plt.figure(figsize=figsize, **kwargs)
        ax = plt.axes()
        ax.set_axisbelow(True)
        plt.grid(color='gainsboro', linewidth=.1)
        plt.title('Target distribution in a single batch', fontweight='bold')
        if self.cat:
            if yb.ndim == 1:
                yb = yb.flatten().detach().cpu().numpy()
                data = np.unique(yb, return_counts=True)[1]
                data = data / np.sum(data)
            else:
                yb = yb.detach().cpu().numpy()
                data = yb.mean(0)
            plt.bar(self.vocab, data, color=color, edgecolor='black')
            plt.xticks(self.vocab)
        else:
            yb = yb.flatten().detach().cpu().numpy()
            weights=np.ones(len(yb)) / len(yb)
            plt.hist(yb, bins=min(len(yb) // 2, 100), weights=weights, color='violet', edgecolor='black')
        plt.gca().yaxis.set_major_formatter(PercentFormatter(1))
        plt.show()

    @property
    def c(self):
        if len(self.dataset) == 0: return 0
        elif hasattr(self, "vocab"):
            return len(self.vocab)
        else:
            return 1
#             return self.d if not is_listy(self.d) else reduce(lambda x, y: x * y, self.d, 1)

    @property
    def d(self):
        if len(self.dataset) == 0: return None
        b = self.one_batch()
        if len(b) == 1: return None
        i = getattr(self, 'n_inp', len(b)-1)
        yb = b[i:]
        if len(yb[0].shape[1:]) == 0: return None
        elif len(yb[0].shape[1:]) == 1: return yb[0].shape[-1]
        else: return list(yb[0].shape[1:])

    @property
    def cat(self): return hasattr(self, "vocab")

    @property
    def cws(self):
        if self.cat:
            if self.ptls[-1].ndim == 2: # one hot encoded
                pos_per_class = self.ptls[-1].sum(0)
                neg_per_class = (len(self.ptls[-1]) - pos_per_class)
                pos_weights = neg_per_class / pos_per_class
                pos_weights[pos_weights == float('inf')] = 0
                return torch.Tensor(pos_weights).to(self.device)
            else:
                counts = torch.unique(self.ptls[-1].flatten(), return_counts=True, sorted=True)[-1]
                iw = (counts.sum() / counts)
                return (iw / iw.sum()).to(self.device)
        else: return None

    @property
    def class_priors(self):
        if self.cws is not None:
            cp = 1. / (self.cws + 1e-8)
            return (cp / cp.sum()).to(self.device)
        else: return None


class TSDataLoader(NumpyDataLoader):
    @property
    def vars(self):
        if len(self.dataset) == 0: return 0
        b = self.one_batch()
        i = getattr(self, 'n_inp', 1 if len(b)==1 else len(b)-1)
        xb = b[:i]
        if hasattr(xb[0], 'vars'): return xb[0].vars
        if xb[0].ndim >= 4: return xb[0].shape[-3]
        else: return xb[0].shape[-2]
    @property
    def len(self):
        if len(self.dataset) == 0: return 0
        b = self.one_batch()
        i = getattr(self, 'n_inp', 1 if len(b)==1 else len(b)-1)
        xb = b[:i]
        if hasattr(xb[0], 'len'): return xb[0].len
        if xb[0].ndim >= 4: return xb[0].shape[-2:]
        else: return xb[0].shape[-1]

# %% ../../nbs/013_data.core.ipynb 62
_batch_tfms = ('after_item','before_batch','after_batch')

class NumpyDataLoaders(DataLoaders):
    _xblock = NumpyTensorBlock
    _dl_type = NumpyDataLoader
    def __init__(self, *loaders, path='.', device=None):
        create_dir(path, verbose=False)
        self.loaders, self.path = list(loaders), Path(path)
        self.device = ifnone(device, default_device())

    def new_dl(self, X, y=None, bs=64):
        assert X.ndim == 3, "You must pass an X with 3 dimensions [batch_size x n_vars x seq_len]"
        if y is not None and not is_array(y) and not is_listy(y): y = [y]
        new_dloader = self.new(self.dataset.add_dataset(X, y=y), bs=min(bs, len(X)))
        return new_dloader

    @delegates(plt.subplots)
    def show_dist(self, figsize=None, **kwargs): self.loaders[0].show_dist(figsize=figsize, **kwargs)

    def decoder(self, o):
        if isinstance(o, tuple): return self.decode(o)
        if o.ndim <= 1: return self.decodes(o)
        else: return L([self.decodes(oi) for oi in o])

    def decode(self, b):
        return to_cpu(self.after_batch.decode(self._retain_dl(b)))


    @classmethod
    @delegates(DataLoaders.from_dblock)
    def from_numpy(cls, X, y=None, splitter=None, valid_pct=0.2, seed=0, item_tfms=None, batch_tfms=None, **kwargs):
        "Create timeseries dataloaders from arrays (X and y, unless unlabeled)"
        if splitter is None: splitter = RandomSplitter(valid_pct=valid_pct, seed=seed)
        getters = [ItemGetter(0), ItemGetter(1)] if y is not None else [ItemGetter(0)]
        dblock = DataBlock(blocks=(cls._xblock, CategoryBlock),
                           getters=getters,
                           splitter=splitter,
                           item_tfms=item_tfms,
                           batch_tfms=batch_tfms)

        source = itemify(X) if y is None else itemify(X,y)
        return cls.from_dblock(dblock, source, **kwargs)

    @classmethod
    def from_dsets(cls, *ds, path='.', bs=64, num_workers=0, batch_tfms=None, device=None, shuffle_train=True, drop_last=True, 
                   weights=None, partial_n=None, sampler=None, sort=False, **kwargs):
        device = ifnone(device, default_device())
        if batch_tfms is not None and not isinstance(batch_tfms, list): batch_tfms = [batch_tfms]
        shuffle_default = (shuffle_train,) + (False,) * (len(ds)-1)
        drop_last_default = (drop_last,) + (False,) * (len(ds)-1)
        defaults = {'shuffle': shuffle_default, 'drop_last': drop_last_default}
        kwargs = merge(defaults, {k: tuplify(v, match=ds) for k,v in kwargs.items()})
        kwargs = [{k: v[i] for k,v in kwargs.items()} for i in range_of(ds)]
        if not is_listy(bs): bs = [bs]
        if len(bs) != len(ds): bs = bs * len(ds)
        if weights is None: weights = [None] * len(ds)
        if not is_listy(partial_n): partial_n = [partial_n]
        if len(partial_n) != len(ds): partial_n = partial_n * len(ds)
        if not is_listy(sampler): sampler = [sampler]
        if len(sampler) != len(ds): sampler = sampler * len(ds)
        loaders = [cls._dl_type(d, bs=b, num_workers=num_workers, batch_tfms=batch_tfms, weights=w, partial_n=n, sampler=s, sort=sort, **k)\
                   for d,k,b,w,n,s in zip(ds, kwargs, bs, weights, partial_n, sampler)]
        return cls(*loaders, path=path, device=device)


class TSDataLoaders(NumpyDataLoaders):
    _xblock = TSTensorBlock
    _dl_type = TSDataLoader

# %% ../../nbs/013_data.core.ipynb 63
class StratifiedSampler:
    "Sampler where batches preserve the percentage of samples for each class"
    
    def __init__(self, 
        y, # The target variable for supervised learning problems. Stratification is done based on the y labels.
        bs : int = 64, # Batch size
        shuffle : bool = False, # Flag to shuffle each class’s samples before splitting into batches.
        drop_last : bool = False # Flag to drop the last incomplete batch.
        ):
        self.n_splits = len(y) // bs if drop_last else int(np.ceil(len(y) / bs))
        self.skf = StratifiedKFold(n_splits=self.n_splits, shuffle=shuffle)
        self.y = y
        self.shuffle = shuffle
        self.n = self.n_splits * bs if drop_last else len(y)

    def __iter__(self):
        if self.shuffle:
            self.skf.random_state = np.random.randint(0, 1e8)
        chunklist = []
        for _,idx in self.skf.split(self.y, self.y):
            chunklist.extend(idx)
        yield chunklist

    def __len__(self):
        return self.n

# %% ../../nbs/013_data.core.ipynb 65
def get_c(dls):
    if getattr(dls, 'c', False): return dls.c
    if getattr(getattr(dls.train, 'after_item', None), 'c', False): return dls.train.after_item.c
    if getattr(getattr(dls.train, 'after_batch', None), 'c', False): return dls.train.after_batch.c
    vocab = getattr(dls, 'vocab', [])
    if len(vocab) > 0 and is_listy(vocab[-1]): vocab = vocab[-1]
    return len(vocab)

# %% ../../nbs/013_data.core.ipynb 66
def get_best_dl_params(dl, n_iters=10, num_workers=[0, 1, 2, 4, 8], pin_memory=[True, False], prefetch_factor=[2, 4, 8], return_best=True, 
                       verbose=True): 

    if not torch.cuda.is_available(): 
        num_workers = 0
    n_iters = min(n_iters, len(dl))
    if not return_best: verbose = True
    
    nw = dl.fake_l.num_workers
    pm = dl.fake_l.pin_memory
    pf = dl.fake_l.prefetch_factor

    try:
        best_nw = nw
        best_pm = pm
        best_pf = pf
        
        # num_workers
        if not num_workers: best_nw = nw
        elif isinstance(num_workers, Integral): best_nw = num_workers
        else: 
            best_time = np.inf
            for _nw in num_workers:
                dl.fake_l.num_workers = _nw
                timer.start(False)
                for i, _ in enumerate(dl):
                    if i == n_iters - 1: 
                        t = timer.stop() / (i + 1)
                        pv(f'   num_workers: {_nw:2}  pin_memory: {pm!s:^5}  prefetch_factor: {pf:2}  -  time: {1_000 * t/n_iters:8.3f} ms/iter', verbose)
                        if t < best_time: 
                            best_nw = _nw
                            best_time = t
                        break
        dl.fake_l.num_workers = best_nw
        

        # pin_memory
        if not pin_memory: best_pm = pm
        elif isinstance(pin_memory, bool): best_pm = pin_memory
        else:
            best_time = np.inf
            if not pin_memory: pin_memory = [pm]
            for _pm in pin_memory:
                dl.fake_l.pin_memory = _pm
                timer.start(False)
                for i, _ in enumerate(dl):
                    if i == n_iters - 1: 
                        t = timer.stop() / (i + 1)
                        pv(f'   num_workers: {best_nw:2}  pin_memory: {_pm!s:^5}  prefetch_factor: {pf:2}  -  time: {1_000 * t/n_iters:8.3f} ms/iter', 
                           verbose)
                        if t < best_time: 
                            best_pm = _pm
                            best_time = t
                        break
        dl.fake_l.pin_memory = best_pm

        # prefetch_factor
        if best_nw == 0: best_pf = 2
        elif not prefetch_factor: best_pf = pf
        elif isinstance(prefetch_factor, Integral): best_pf = prefetch_factor
        else:
            best_time = np.inf
            if not prefetch_factor: prefetch_factor = [pf]
            for _pf in prefetch_factor:
                dl.fake_l.prefetch_factor = _pf
                timer.start(False)
                for i, _ in enumerate(dl):
                    if i == n_iters - 1: 
                        t = timer.stop() / (i + 1)
                        pv(f'   num_workers: {best_nw:2}  pin_memory: {best_pm!s:^5}  prefetch_factor: {_pf:2}  -  time: {1_000 * t/n_iters:8.3f} ms/iter', 
                           verbose)
                        if t < best_time: 
                            best_pf = _pf
                            best_time = t
                        break
        dl.fake_l.prefetch_factor = best_pf
    
    except KeyboardInterrupt: 
        dl.fake_l.num_workers = best_nw if return_best else nw
        dl.fake_l.pin_memory = best_pm if return_best else pm
        dl.fake_l.prefetch_factor = best_pf if return_best else pf

    if not return_best: 
        dl.fake_l.num_workers = nw
        dl.fake_l.pin_memory = pm
        dl.fake_l.prefetch_factor = pf

    if verbose: 
        print('\n   best dl params:')
        print(f'       best num_workers    : {best_nw}')
        print(f'       best pin_memory     : {best_pm}')
        print(f'       best prefetch_factor: {best_pf}')
        print(f'       return_best         : {return_best}')
        print('\n')

    return dl

def get_best_dls_params(dls, n_iters=10, num_workers=[0, 1, 2, 4, 8], pin_memory=[True, False], prefetch_factor=[2, 4, 8], 
                        return_best=True, verbose=True): 
    
    for i in range(len(dls.loaders)):
        try:
            print(f'\nDataloader {i}\n')
            dls.loaders[i] = get_best_dl_params(dls.loaders[i], n_iters=n_iters, num_workers=num_workers, pin_memory=pin_memory, 
                                            prefetch_factor=prefetch_factor, return_best=return_best, verbose=verbose)
        except KeyboardInterrupt: pass
    return dls

# %% ../../nbs/013_data.core.ipynb 67
def get_ts_dls(X, y=None, splits=None, sel_vars=None, sel_steps=None, tfms=None, inplace=True,
               path='.', bs=64, batch_tfms=None, num_workers=0, device=None, shuffle_train=True, drop_last=True, 
               weights=None, partial_n=None, sampler=None, sort=False, **kwargs):
    if splits is None: splits = (L(np.arange(len(X)).tolist()), L([]))
    create_dir(path, verbose=False)
    dsets = TSDatasets(X, y, splits=splits, sel_vars=sel_vars, sel_steps=sel_steps, tfms=tfms, inplace=inplace)
    if weights is not None:
        assert len(X) == len(weights)
        if splits is not None: 
            weights = [weights[split] if i == 0 else None for i,split in enumerate(splits)] # weights only applied to train set
    dls   = TSDataLoaders.from_dsets(dsets.train, dsets.valid, path=path, bs=bs, batch_tfms=batch_tfms, num_workers=num_workers,
                                     device=device, shuffle_train=shuffle_train, drop_last=drop_last, weights=weights, 
                                     partial_n=partial_n, sampler=sampler, sort=sort, **kwargs)
    return dls

def get_ts_dl(X, y=None, split=None, sel_vars=None, sel_steps=None, tfms=None, inplace=True,
              path='.', bs=64, batch_tfms=None, num_workers=0, device=None, shuffle_train=True, drop_last=True, weights=None,
              partial_n=None, sampler=None, sort=False, **kwargs):
    if split is None: split = L(np.arange(len(X)).tolist())
    splits = (split, L([]))
    create_dir(path, verbose=False)
    dsets = TSDatasets(X, y, splits=splits, sel_vars=sel_vars, sel_steps=sel_steps, tfms=tfms, inplace=inplace, **kwargs)
    if not is_listy(partial_n): partial_n = [partial_n]
    dls   = TSDataLoaders.from_dsets(dsets.train, path=path, bs=bs, batch_tfms=batch_tfms, num_workers=num_workers,
                                     device=device, shuffle_train=shuffle_train, drop_last=drop_last, weights=weights,
                                     partial_n=partial_n, sampler=sampler, sort=sort, **kwargs)
    return dls.train

get_tsimage_dls = get_ts_dls

def get_subset_dl(dl, idxs): return dl.new(dl.dataset.subset(idxs))

# %% ../../nbs/013_data.core.ipynb 105
def get_time_per_batch(dl, model=None, n_batches=None):
    try:
        timer.start(False)
        pbar = progress_bar(dl, leave=False)
        for i, (xb, _) in enumerate(pbar): 
            if model is not None: 
                _ = model(xb)
            if n_batches is not None and i >= n_batches - 1: 
                t = timer.stop()
                pbar.on_interrupt()
                break
        if n_batches is None or i < n_batches - 1: 
            t = timer.stop()
        
    except KeyboardInterrupt:
        t = timer.stop()
        pbar.on_interrupt()
    return t / (i+1)

def get_dl_percent_per_epoch(dl, model, n_batches=None): 
    dl_time = get_time_per_batch(dl, model=None, n_batches=n_batches)
    model_time = get_time_per_batch(dl, model=model, n_batches=n_batches)
    return f'{min(1, dl_time/model_time):.2%}'
