Source code for sklego.datasets

import os
import numpy as np
import pandas as pd
from pkg_resources import resource_filename

from sklearn.datasets import fetch_openml


[docs]def load_penguins(return_X_y=False, as_frame=False): """ Loads the penguins dataset, which is a lovely alternative for the iris dataset. We've added this dataset for educational use. Data were collected and made available by Dr. Kristen Gorman and the Palmer Station, Antarctica LTER, a member of the Long Term Ecological Research Network. The goal of the dataset is to predict which species of penguin a penguin belongs to. This data originally appeared as a R package and R users can find this data in the palmerpenguins package https://github.com/allisonhorst/palmerpenguins. You can also go to the repository for some lovely images that explain the dataset. To cite this dataset in publications use: Gorman KB, Williams TD, Fraser WR (2014) Ecological Sexual Dimorphism and Environmental Variability within a Community of Antarctic Penguins (Genus Pygoscelis). PLoS ONE 9(3): e90081. https://doi.org/10.1371/journal.pone.0090081 :param return_X_y: If True, returns ``(data, target)`` instead of a dict object. :param as_frame: give the pandas dataframe instead of X, y matrices (default=False) :Example: >>> from sklego.datasets import load_penguins >>> X, y = load_penguins(return_X_y=True) >>> X.shape (344, 6) >>> y.shape (344,) >>> load_penguins(as_frame=True).columns Index(['species', 'island', 'bill_length_mm', 'bill_depth_mm', 'flipper_length_mm', 'body_mass_g', 'sex'], dtype='object') Additional data use information ############################### Anyone interested in publishing the data should contact `Dr. Kristen Gorman <https://www.uaf.edu/cfos/people/faculty/detail/kristen-gorman.php>`_ about analysis and working together on any final products. From Gorman et al. (2014): > “Data reported here are publicly available within the PAL-LTER data > system (datasets 219, 220, and 221): > `<http://oceaninformatics.ucsd.edu/datazoo/data/pallter/datasets>`_. > Individuals interested in using these data are therefore expected to > follow the US LTER Network’s Data Access Policy, Requirements and Use > Agreement: `<https://lternet.edu/data-access-policy/>`_.” **Please cite data using the following:** **Adélie penguins:** - Palmer Station Antarctica LTER and K. Gorman, 2020. Structural size measurements and isotopic signatures of foraging among adult male and female Adélie penguins (*Pygoscelis adeliae*) nesting along the Palmer Archipelago near Palmer Station, 2007-2009 ver 5. Environmental Data Initiative. `<https://doi.org/10.6073/pasta/98b16d7d563f265cb52372c8ca99e60f>`_ (Accessed 2020-06-08). **Gentoo penguins:** - Palmer Station Antarctica LTER and K. Gorman, 2020. Structural size measurements and isotopic signatures of foraging among adult male and female Gentoo penguin (*Pygoscelis papua*) nesting along the Palmer Archipelago near Palmer Station, 2007-2009 ver 5. Environmental Data Initiative. `<https://doi.org/10.6073/pasta/7fca67fb28d56ee2ffa3d9370ebda689>`_ (Accessed 2020-06-08). **Chinstrap penguins:** - Palmer Station Antarctica LTER and K. Gorman, 2020. Structural size measurements and isotopic signatures of foraging among adult male and female Chinstrap penguin (*Pygoscelis antarcticus*) nesting along the Palmer Archipelago near Palmer Station, 2007-2009 ver 6. Environmental Data Initiative. `<https://doi.org/10.6073/pasta/c14dfcfada8ea13a17536e73eb6fbe9e>`_ (Accessed 2020-06-08). """ filepath = resource_filename("sklego", os.path.join("data", "penguins.zip")) df = pd.read_csv(filepath) if as_frame: return df X, y = ( df[ [ "island", "bill_length_mm", "bill_depth_mm", "flipper_length_mm", "body_mass_g", "sex", ] ].values, df["species"].values, ) if return_X_y: return X, y return {"data": X, "target": y}
[docs]def load_arrests(return_X_y=False, as_frame=False): """ Loads the arrests dataset which can serve as a benchmark for fairness. It is data on the police treatment of individuals arrested in Toronto for simple possession of small quantities of marijuana. The goal is to predict whether or not the arrestee was released with a summons while maintaining a degree of fairness. :param return_X_y: If True, returns ``(data, target)`` instead of a dict object. :param as_frame: give the pandas dataframe instead of X, y matrices (default=False) :Example: >>> from sklego.datasets import load_arrests >>> X, y = load_arrests(return_X_y=True) >>> X.shape (5226, 7) >>> y.shape (5226,) >>> load_arrests(as_frame=True).columns Index(['released', 'colour', 'year', 'age', 'sex', 'employed', 'citizen', 'checks'], dtype='object') The dataset was copied from the carData R package and can originally be found in: - Personal communication from Michael Friendly, York University. The documentation page of the dataset from the package can be viewed here: http://vincentarelbundock.github.io/Rdatasets/doc/carData/Arrests.html """ filepath = resource_filename("sklego", os.path.join("data", "arrests.zip")) df = pd.read_csv(filepath) if as_frame: return df X, y = ( df[["colour", "year", "age", "sex", "employed", "citizen", "checks"]].values, df["released"].values, ) if return_X_y: return X, y return {"data": X, "target": y}
[docs]def load_chicken(return_X_y=False, as_frame=False): """ Loads the chicken dataset. The chicken data has 578 rows and 4 columns from an experiment on the effect of diet on early growth of chicks. The body weights of the chicks were measured at birth and every second day thereafter until day 20. They were also measured on day 21. There were four groups on chicks on different protein diets. :param return_X_y: If True, returns ``(data, target)`` instead of a dict object. :param as_frame: give the pandas dataframe instead of X, y matrices (default=False) :Example: >>> from sklego.datasets import load_chicken >>> X, y = load_chicken(return_X_y=True) >>> X.shape (578, 3) >>> y.shape (578,) >>> load_chicken(as_frame=True).columns Index(['weight', 'time', 'chick', 'diet'], dtype='object') The datasets can be found in the following sources: - Crowder, M. and Hand, D. (1990), Analysis of Repeated Measures, Chapman and Hall (example 5.3) - Hand, D. and Crowder, M. (1996), Practical Longitudinal Data Analysis, Chapman and Hall (table A.2) """ filepath = resource_filename("sklego", os.path.join("data", "chickweight.zip")) df = pd.read_csv(filepath) if as_frame: return df X, y = df[["time", "diet", "chick"]].values, df["weight"].values if return_X_y: return X, y return {"data": X, "target": y}
[docs]def load_abalone(return_X_y=False, as_frame=False): """ Loads the abalone dataset where the goal is to predict the gender of the creature. :param return_X_y: If True, returns ``(data, target)`` instead of a dict object. :param as_frame: give the pandas dataframe instead of X, y matrices (default=False) :Example: >>> from sklego.datasets import load_abalone >>> X, y = load_abalone(return_X_y=True) >>> X.shape (4177, 8) >>> y.shape (4177,) >>> load_abalone(as_frame=True).columns Index(['sex', 'length', 'diameter', 'height', 'whole_weight', 'shucked_weight', 'viscera_weight', 'shell_weight', 'rings'], dtype='object') The dataset was copied from Kaggle and can originally be found in: can be found in the following sources: - Warwick J Nash, Tracy L Sellers, Simon R Talbot, Andrew J Cawthorn and Wes B Ford (1994) "The Population Biology of Abalone (_Haliotis_ species) in Tasmania." Sea Fisheries Division, Technical Report No. 48 (ISSN 1034-3288) """ filepath = resource_filename("sklego", os.path.join("data", "abalone.zip")) df = pd.read_csv(filepath) if as_frame: return df X = df[ [ "length", "diameter", "height", "whole_weight", "shucked_weight", "viscera_weight", "shell_weight", "rings", ] ].values y = df["sex"].values if return_X_y: return X, y return {"data": X, "target": y}
[docs]def load_heroes(return_X_y=False, as_frame=False): """ A dataset from a video game: "heroes of the storm". The goal of the dataset is to predict the attack type. Note that the pandas dataset returns more information. This is because we wanted to keep the X simple in the return_X_y case. :param return_X_y: If True, returns ``(data, target)`` instead of a dict object. :param as_frame: give the pandas dataframe instead of X, y matrices (default=False) :Example: >>> X, y = load_heroes(return_X_y=True) >>> X.shape (84, 2) >>> y.shape (84,) >>> df = load_heroes(as_frame=True) >>> df.columns Index(['name', 'attack_type', 'role', 'health', 'attack', 'attack_spd'], dtype='object') """ filepath = resource_filename("sklego", os.path.join("data", "heroes.zip")) df = pd.read_csv(filepath) if as_frame: return df X = df[["health", "attack"]].values y = df["attack_type"].values if return_X_y: return X, y return {"data": X, "target": y}
[docs]def make_simpleseries( n_samples=365 * 5, trend=0.001, season_trend=0.001, noise=0.5, as_frame=False, seed=None, stack_noise=False, start_date=None, ): """ Generate a very simple timeseries dataset to play with. The generator returns a daily time-series with a yearly seasonality, trend, and noise. :param n_samples: The number of days to simulate the timeseries for. :param trend: The long term trend in the dataset. :param season_trend: The long term trend in the seasonality. :param noise: The noise that is applied to the dataset. :param as_frame: Return a pandas dataframe instead of a numpy array. :param seed: The seed value for the randomness. :param stack_noise: Set the noise to be stacked by a cumulative sum. :param start_date: Also add a start date (only works if `as_frame`=True). :return: numpy array unless dataframe is specified :Example: >>> from sklego.datasets import make_simpleseries >>> make_simpleseries(seed=42) array([-0.34078806, -0.61828731, -0.18458236, ..., -0.27547402, -0.38237413, 0.13489355]) >>> make_simpleseries(as_frame=True, start_date="2018-01-01", seed=42).head(3) yt date 0 -0.340788 2018-01-01 1 -0.618287 2018-01-02 2 -0.184582 2018-01-03 """ if seed: np.random.seed(seed) time = np.arange(0, n_samples) noise = np.random.normal(0, noise, n_samples) if stack_noise: noise = noise.cumsum() r1, r2 = np.random.normal(0, 1, 2) seasonality = r1 * np.sin(time / 365 * 2 * np.pi) + r2 * np.cos( time / 365 * 4 * np.pi + 1 ) result = seasonality + season_trend * seasonality * time + trend * time + noise if as_frame: if start_date: stamps = pd.date_range(start_date, periods=n_samples) return pd.DataFrame({"yt": result, "date": stamps}) return pd.DataFrame({"yt": result}) return result
[docs]def load_hearts(return_X_y=False, as_frame=False): """ Loads the Cleveland Heart Diseases dataset. The goal is to predict the presence of a heart disease (target values 1, 2, 3, and 4). The data originates from research to heart diseases by four institutions and originally contains 76 attributes. Yet, all published experiments refer to using a subset of 13 features and one target. This implementation loads the Cleveland dataset of the research which is the only set used by ML researchers to this date. :param return_X_y: If True, returns ``(data, target)`` instead of a dict object. :param as_frame: give the pandas dataframe instead of X, y matrices (default=False) :Example: >>> X, y = load_hearts(return_X_y=True) >>> X.shape (303, 13) >>> y.shape (303,) >>> df = load_hearts(as_frame=True) >>> df.columns Index(['age', 'sex', 'cp', 'trestbps', 'chol', 'fbs', 'restecg', 'thalach', 'exang', 'oldpeak', 'slope', 'ca', 'thal', 'target'], dtype='object') The dataset can originally be found here: https://archive.ics.uci.edu/ml/datasets/Heart+Disease The responsible institutions for the entire database are: 1. Hungarian Institute of Cardiology. Budapest: Andras Janosi, M.D. 2. University Hospital, Zurich, Switzerland: William Steinbrunn, M.D. 3. University Hospital, Basel, Switzerland: Matthias Pfisterer, M.D. 4. V.A. Medical Center, Long Beach and Cleveland Clinic Foundation: Robert Detrano, M.D., Ph.D. The documentation of the dataset can be viewed here: https://archive.ics.uci.edu/ml/machine-learning-databases/heart-disease/heart-disease.names """ filepath = resource_filename("sklego", os.path.join("data", "hearts.zip")) df = pd.read_csv(filepath) if as_frame: return df X = df[ [ "age", "sex", "cp", "trestbps", "chol", "fbs", "restecg", "thalach", "exang", "oldpeak", "slope", "ca", "thal", ] ].values y = df["target"].values if return_X_y: return X, y return {"data": X, "target": y}
[docs]def fetch_creditcard(*, cache=True, data_home=None, as_frame=False, return_X_y=False): """ Load the creditcard dataset. Download it if necessary. Note that internally this is using `fetch_openml` from scikit-learn, which is experimental. ============== ============== Samples total 284807 Dimensionality 29 Features real Target int 0, 1 ============== ============== The datasets contains transactions made by credit cards in September 2013 by european cardholders. This dataset present transactions that occurred in two days, where we have 492 frauds out of 284,807 transactions. The dataset is highly unbalanced, the positive class (frauds) account for 0.172% of all transactions. Please cite: Andrea Dal Pozzolo, Olivier Caelen, Reid A. Johnson and Gianluca Bontempi. Calibrating Probability with Undersampling for Unbalanced Classification. In Symposium on Computational Intelligence and Data Mining (CIDM), IEEE, 2015 :param version: integer or 'active', default='active' Version of the dataset. Can only be provided if also ``name`` is given. If 'active' the oldest version that's still active is used. Since there may be more than one active version of a dataset, and those versions may fundamentally be different from one another, setting an exact version is highly recommended. :param cache: boolean, default=True Whether to cache downloaded datasets using joblib. :param data_home: optional, default: None Specify another download and cache folder for the datasets. By default all scikit-learn data is stored in '~/scikit_learn_data' subfolders. :param as_frame: boolean, default=False If True, the data is a pandas DataFrame including columns with appropriate dtypes (numeric, string or categorical). The target is a pandas DataFrame or Series depending on the number of target_columns. The Bunch will contain a ``frame`` attribute with the target and the data. If ``return_X_y`` is True, then ``(data, target)`` will be pandas DataFrames or Series as describe above. :param return_X_y: : boolean, default=False. If True, returns ``(data.data, data.target)`` instead of a Bunch object. Returns ------- :return: Dictionary-like object, with the following attributes. * data ndarray, shape (284807, 29) if ``as_frame`` is True, ``data`` is a pandas object. * target ndarray, shape (284807, ) if ``as_frame`` is True, ``target`` is a pandas object. * feature_names Array of ordered feature names used in the dataset. * DESCR Description of the creditcard dataset. Best to use print. Notes ----- This dataset consists of 284807 samples and 29 features. """ return fetch_openml( data_id=1597, data_home=data_home, cache=cache, as_frame=as_frame, return_X_y=return_X_y, )