doubt.datasets.dataset

  1"""Base class for data sets"""
  2
  3import re
  4import warnings
  5from abc import ABC, abstractmethod
  6from typing import Iterable, Optional, Tuple, Union
  7
  8import numpy as np
  9import pandas as pd
 10import requests
 11
 12BASE_DATASET_DESCRIPTION = """
 13    Parameters:
 14        cache (str or None, optional):
 15            The name of the cache. It will be saved to `cache` in the current working
 16            directory. If None then no cache will be saved. Defaults to
 17            '.dataset_cache'.
 18
 19    Attributes:
 20        cache (str or None):
 21            The name of the cache.
 22        shape (tuple of integers):
 23            Dimensions of the data set
 24        columns (list of strings):
 25            List of column names in the data set
 26"""
 27
 28
 29class BaseDataset(ABC):
 30    _url: str
 31    _features: Iterable
 32    _targets: Iterable
 33
 34    def __init__(self, cache: Optional[str] = ".dataset_cache"):
 35        self.cache = cache
 36        self._data = self.get_data()
 37        self.shape = self._data.shape
 38        self.columns = self._data.columns
 39
 40    @abstractmethod
 41    def _prep_data(self, data: bytes) -> pd.DataFrame:
 42        return
 43
 44    def get_data(self) -> pd.DataFrame:
 45        """Download and prepare the dataset.
 46
 47        Returns:
 48            Pandas DataFrame: The dataset.
 49        """
 50
 51        # Get name of dataset, being the class name converted to snake case
 52        name = re.sub(r"([A-Z])", r"_\1", type(self).__name__)
 53        name = name.lower().strip("_")
 54
 55        try:
 56            if self.cache is not None:
 57                data = pd.read_hdf(self.cache, name)
 58        except (FileNotFoundError, KeyError):
 59            with warnings.catch_warnings():
 60                warnings.simplefilter("ignore")
 61                response = requests.get(self._url, verify=False)
 62            data = self._prep_data(response.content)
 63            if self.cache is not None:
 64                data.to_hdf(self.cache, name)
 65        return data
 66
 67    def to_pandas(self) -> pd.DataFrame:
 68        return self._data
 69
 70    def __len__(self) -> int:
 71        return len(self._data)
 72
 73    def head(self, n: int = 5) -> pd.DataFrame:
 74        return self._data.head(n)
 75
 76    def close(self):
 77        del self._data
 78        del self
 79
 80    def __exit__(self, exc_type: str, exc_value: str, exc_traceback: str):
 81        self.close()
 82
 83    def __str__(self) -> str:
 84        return str(self._data)
 85
 86    def __repr__(self) -> str:
 87        return repr(self._data)
 88
 89    def _repr_html_(self):
 90        return self._data._repr_html_()
 91
 92    def split(
 93        self, test_size: Optional[float] = None, random_seed: Optional[float] = None
 94    ) -> Union[
 95        Tuple[np.ndarray, np.ndarray],
 96        Tuple[np.ndarray, np.ndarray, np.ndarray, np.ndarray],
 97    ]:
 98        """Split dataset into features and targets and optionally train/test.
 99
100        Args:
101            test_size (float or None):
102                The fraction of the dataset that will constitute the test set. If None
103                then no train/test split will happen. Defaults to None.
104            random_seed (float or None):
105                The random seed used for the train/test split. If None then a random
106                number will be chosen. Defaults to None.
107
108        Returns:
109            If `test_size` is not `None` then a tuple of numpy arrays
110            (X_train, y_train, X_test, y_test) is returned, and otherwise
111            the tuple (X, y) of numpy arrays is returned.
112        """
113        # Initialise random number generator
114        rng = np.random.default_rng(random_seed)
115
116        nrows = len(self._data)
117        features = self._features
118        targets = self._targets
119
120        if test_size is not None:
121            test_idxs = rng.random(size=(nrows,)) < test_size
122            train_idxs = ~test_idxs
123
124            X_train = self._data.iloc[train_idxs, features].values
125            y_train = self._data.iloc[train_idxs, targets].values.squeeze()
126            X_test = self._data.iloc[test_idxs, features].values
127            y_test = self._data.iloc[test_idxs, targets].values.squeeze()
128
129            return X_train, X_test, y_train, y_test
130
131        else:
132            X = self._data.iloc[:, features].values
133            y = self._data.iloc[:, targets].values.squeeze()
134            return X, y