Skip to content

Datasets

Datasets

Dataset

Simple container for train/test arrays and per-group indices.

__init__(name: str, X_train: np.ndarray, X_test: np.ndarray, y_train: np.ndarray, y_test: np.ndarray, group_idx: list, group_idx_test: list)

Parameters:

  • name (str) –

    Dataset name.

  • X_train (ndarray) –

    Feature matrices.

  • X_test (ndarray) –

    Feature matrices.

  • y_train (ndarray) –

    Targets.

  • y_test (ndarray) –

    Targets.

  • group_idx (list) –

    Per-group selectors for train/test. Each entry is either a boolean mask over the corresponding split or an array/list of integer indices.

  • group_idx_test (list) –

    Per-group selectors for train/test. Each entry is either a boolean mask over the corresponding split or an array/list of integer indices.

subsample(n: int, random_state: int | None = 42) -> Dataset

Subsample the training split and rebuild training group indices.

Parameters:

  • n (int) –

    Number of training samples to keep.

  • random_state (int | None, default: 42 ) –

    Seed for sampling. If None, uses an unseeded generator.

Returns:

  • Dataset –

    New dataset with subsampled training data and updated groups.

Raises:

  • ValueError –

    If n is not in [1, n_train], or if a boolean group mask has the wrong length.

fetch_ACSEmployment(states: List[str] = ['WY'], year: int = 2018, n_groups: int = 2, test_size: float = 0.3, random_state: int = 42) -> Dataset

Fetch Folktables ACSEmployment (ESR == 1) and build groups from SEX × RAC1P.

Parameters:

  • states (list[str], default: ["WY"] ) –

    State FIPS codes accepted by folktables.

  • year ((2014, 2015, 2016, 2017, 2018), default: 2014 ) –

    ACS 1-year survey release year.

  • n_groups (int, default: 2 ) –

    Number of groups to return.

  • test_size (float, default: 0.3 ) –

    Test split size.

  • random_state (int, default: 42 ) –

    Train/test split seed.

Returns:

  • Dataset –

    Train/test splits and per-group indices.

fetch_ACSIncome(states: List[str] = ['WY'], year: int = 2018, n_groups: int = 2, test_size: float = 0.3, random_state: int = 42) -> Dataset

Fetch Folktables ACSIncome (binary income > 50k) and build groups from SEX × RAC1P.

Parameters:

  • states (list[str], default: ["WY"] ) –

    State FIPS codes accepted by folktables.

  • year ((2014, 2015, 2016, 2017, 2018), default: 2014 ) –

    ACS 1-year survey release year.

  • n_groups (int, default: 2 ) –

    Number of groups to return.

  • test_size (float, default: 0.3 ) –

    Test split size.

  • random_state (int, default: 42 ) –

    Train/test split seed.

Returns:

  • Dataset –

    Train/test splits and per-group indices.

fetch_ACSIncomeR(states: List[str] = ['WY'], year: int = 2018, n_groups: int = 2, test_size: float = 0.3, random_state: int = 42) -> Dataset

Fetch Folktables ACSIncomeR (PINCP regression) and build groups from SEX × RAC1P.

Parameters:

  • states (list[str], default: ["WY"] ) –

    State FIPS codes accepted by folktables.

  • year ((2014, 2015, 2016, 2017, 2018), default: 2014 ) –

    ACS 1-year survey release year.

  • n_groups (int, default: 2 ) –

    Number of groups to return.

  • test_size (float, default: 0.3 ) –

    Test split size.

  • random_state (int, default: 42 ) –

    Train/test split seed.

Returns:

  • Dataset –

    Train/test splits and per-group indices.

fetch_ACSTravelTime(states: List[str] = ['WY'], year: int = 2018, n_groups: int = 2, test_size: float = 0.3, random_state: int = 42) -> Dataset

Fetch Folktables ACSTravelTime (JWMNP > 20) and build groups from SEX × RAC1P × AGEP.

Parameters:

  • states (list[str], default: ["WY"] ) –

    State FIPS codes accepted by folktables.

  • year ((2014, 2015, 2016, 2017, 2018), default: 2014 ) –

    ACS 1-year survey release year.

  • n_groups (int, default: 2 ) –

    Number of groups to return.

  • test_size (float, default: 0.3 ) –

    Test split size.

  • random_state (int, default: 42 ) –

    Train/test split seed.

Returns:

  • Dataset –

    Train/test splits and per-group indices.

fetch_adult(n_groups: int = 2, test_size: float = 0.3, random_state: int = 42) -> Dataset

Fetch and prepare the Adult dataset (OpenML: adult, v2).

Target: class mapped to {<=50K: 0.0, >50K: 1.0}. Groups: intersections of sex × race (merged until n_groups remain). Numeric features are standardized on the training split.

Parameters:

  • n_groups (int, default: 2 ) –

    Number of groups to return.

  • test_size (float, default: 0.3 ) –

    Test split size.

  • random_state (int, default: 42 ) –

    Train/test split seed.

Returns:

  • Dataset –

    Train/test splits and per-group indices.

Raises:

  • ValueError –

    If n_groups is larger than the number of available (sex, race) groups.

fetch_arrhythmia(n_groups: int = 2, test_size: float = 0.3, random_state: int = 42) -> Dataset

Fetch Arrhythmia (OpenML arrhythmia, v2).

Target: binaryClass -> {N: 0.0, P: 1.0}. Groups: sex (0 vs 1). Only n_groups=2 is supported.

Parameters:

  • n_groups (int, default: 2 ) –

    Must be 2.

  • test_size (float, default: 0.3 ) –

    Test split size.

  • random_state (int, default: 42 ) –

    Train/test split seed.

Returns:

  • Dataset –

    Train/test splits and per-group indices.

Raises:

  • ValueError –

    If n_groups is not 2.

fetch_communities_and_crime(n_groups: int = 2, test_size: float = 0.3, random_state: int = 42) -> Dataset

Fetch Communities and Crime (OpenML us_crime, v1).

Target: ViolentCrimesPerPop (float). Groups: race category from argmax over {black, white, asian, hispanic} race percentage columns, merged until n_groups remain. Features are standardized on the training split.

Parameters:

  • n_groups (int, default: 2 ) –

    Number of groups to return.

  • test_size (float, default: 0.3 ) –

    Test split size.

  • random_state (int, default: 42 ) –

    Train/test split seed.

Returns:

  • Dataset –

    Train/test splits and per-group indices.

Raises:

  • ValueError –

    If n_groups is larger than the number of race categories present.

fetch_compas(n_groups: int = 2, test_size: float = 0.3, random_state: int = 42) -> Dataset

Fetch and prepare COMPAS (ProPublica compas-scores-two-years.csv).

Target: two_year_recid (float in {0.0, 1.0}). Groups: intersections of sex × race (merged until n_groups remain). Features are standardized on the training split.

Parameters:

  • n_groups (int, default: 2 ) –

    Number of groups to return.

  • test_size (float, default: 0.3 ) –

    Test split size.

  • random_state (int, default: 42 ) –

    Train/test split seed.

Returns:

  • Dataset –

    Train/test splits and per-group indices.

Raises:

  • ValueError –

    If n_groups is larger than the number of available (sex, race) groups.

fetch_csv(url: str, data_home: str, filename: str = '', replace: bool = False, **read_csv_kwargs) -> pd.DataFrame

Download a CSV file if needed and return it as a DataFrame.

Parameters:

  • url (str) –

    CSV URL.

  • data_home (str) –

    Directory where the file is cached.

  • filename (str, default: "" ) –

    Local filename. If empty, uses the basename of url.

  • replace (bool, default: False ) –

    If True, re-download even if the file already exists.

  • **read_csv_kwargs –

    Passed to pandas.read_csv.

Returns:

  • DataFrame –

    Loaded CSV.

fetch_germancredit(n_groups: int = 2, test_size: float = 0.3, random_state: int = 42) -> Dataset

Fetch German Credit (UCI Statlog German Credit).

Target: derived from credit (binary). Groups: sex extracted from Personal Status. Only n_groups=2 is supported. Features are standardized on the training split.

Parameters:

  • n_groups (int, default: 2 ) –

    Must be 2.

  • test_size (float, default: 0.3 ) –

    Test split size.

  • random_state (int, default: 42 ) –

    Train/test split seed.

Returns:

  • Dataset –

    Train/test splits and per-group indices.

Raises:

  • ValueError –

    If n_groups is not 2.

fetch_lawschool(n_groups: int = 2, test_size: float = 0.3, random_state: int = 42) -> Dataset

Fetch LSAC (bar pass / law school admissions) dataset.

Target: zgpa. Groups: intersections of gender × race1 (merged until n_groups remain). Features are standardized on the training split.

Parameters:

  • n_groups (int, default: 2 ) –

    Number of groups to return.

  • test_size (float, default: 0.3 ) –

    Test split size.

  • random_state (int, default: 42 ) –

    Train/test split seed.

Returns:

  • Dataset –

    Train/test splits and per-group indices.

Raises:

  • ValueError –

    If n_groups is larger than the number of available (gender, race1) groups.

fetch_parkinsons(n_groups: int = 2, test_size: float = 0.3, random_state: int = 42) -> Dataset

Fetch Parkinson's Telemonitoring (UCI parkinsons_updrs.data).

Target: total_UPDRS. Groups: sex (0 vs 1). Only n_groups=2 is supported. Features are standardized on the training split.

Parameters:

  • n_groups (int, default: 2 ) –

    Must be 2.

  • test_size (float, default: 0.3 ) –

    Test split size.

  • random_state (int, default: 42 ) –

    Train/test split seed.

Returns:

  • Dataset –

    Train/test splits and per-group indices.

Raises:

  • ValueError –

    If n_groups is not 2.

fetch_studentperformance(n_groups: int = 2, test_size: float = 0.3, random_state: int = 42) -> Dataset

Fetch Student Performance (OpenML student-performance-uci, v1).

Target: G3. Groups: sex (M vs F, mapped to 0.0/1.0). Only n_groups=2 is supported. Numeric features are standardized on the training split.

Parameters:

  • n_groups (int, default: 2 ) –

    Must be 2.

  • test_size (float, default: 0.3 ) –

    Test split size.

  • random_state (int, default: 42 ) –

    Train/test split seed.

Returns:

  • Dataset –

    Train/test splits and per-group indices.

Raises:

  • ValueError –

    If n_groups is not 2.

load_dataframe(df: pd.DataFrame, target_col: str, sensitive_cols: Union[str, List[str]], name: str = 'user dataframe', n_groups: int = 2, test_size: float = 0.3, random_state: int = 42) -> Dataset

Build a :class:~badr.datasets.Dataset from a user-provided DataFrame.

Parameters:

  • df (DataFrame) –

    Input data.

  • target_col (str) –

    Target column name.

  • sensitive_cols (str or list[str]) –

    Sensitive column(s) used to form groups. These columns are not included in X.

  • name (str, default: "user dataframe" ) –

    Dataset name stored on the returned Dataset.

  • n_groups (int, default: 2 ) –

    Number of groups to return (by merging sensitive-value buckets if needed).

  • test_size (float, default: 0.3 ) –

    Test split size.

  • random_state (int, default: 42 ) –

    Train/test split seed.

Returns:

  • Dataset –

    Train/test splits and per-group indices.

Raises:

  • ValueError –

    If any sensitive_cols are not present in df, or if n_groups cannot be formed.