examples/parse-folktables-datasets.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "a4d0fb85",
   "metadata": {},
   "source": [
    "# Obtaining parsed folktables datasets"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "5f222039",
   "metadata": {},
   "outputs": [],
   "source": [
    "import sys\n",
    "import logging\n",
    "from pathlib import Path\n",
    "\n",
    "import pandas as pd\n",
    "import numpy as np\n",
    "from folktables import ACSDataSource"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "00d4bd71",
   "metadata": {},
   "source": [
    "**NOTE**: use `MAX_SENSITIVE_GROUPS=2` to generate datasets for binary-group experiments."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "811ad844",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Important constants!\n",
    "TRAIN_SIZE = 0.6\n",
    "TEST_SIZE = 0.2\n",
    "VALIDATION_SIZE = 0.2\n",
    "\n",
    "MAX_SENSITIVE_GROUPS = 4          # keep only samples from the 4 largest groups\n",
    "# MAX_SENSITIVE_GROUPS = None     # keep samples from all groups\n",
    "\n",
    "SEED = 42"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "377fc203",
   "metadata": {},
   "outputs": [],
   "source": [
    "assert TRAIN_SIZE + TEST_SIZE + (VALIDATION_SIZE or 0.) == 1  # sanity check"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "dae0cfba",
   "metadata": {},
   "source": [
    "**Change** these paths according to where you want the data to be saved to."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "60b5f503",
   "metadata": {},
   "outputs": [],
   "source": [
    "root_dir = Path(\"~\").expanduser()\n",
    "data_dir = root_dir / \"data\" / \"folktables\"\n",
    "data_dir.mkdir(parents=True, exist_ok=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "ae28e462",
   "metadata": {},
   "outputs": [],
   "source": [
    "# download 2018 ACS data\n",
    "from folktables.load_acs import state_list\n",
    "\n",
    "data_source = ACSDataSource(\n",
    "    survey_year='2018', horizon='1-Year', survey='person',\n",
    "    root_dir=str(data_dir),\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "8afe4020",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(3236107, 286)"
      ]
     },
     "execution_count": 6,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# data is 3236107 rows x 286 columns\n",
    "acs_data = data_source.get_data(states=state_list, download=True)  # use download=True if not yet downloaded\n",
    "acs_data.shape"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6ea8d039",
   "metadata": {},
   "source": [
    "According to the dataset's datasheet, train/test splits should be stratified by state\n",
    "(at least for ACSIncome, the remaining tasks seem ambiguous)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "c4e23b32",
   "metadata": {},
   "outputs": [],
   "source": [
    "STATE_COL = \"ST\"\n",
    "\n",
    "ACS_CATEGORICAL_COLS = {\n",
    "    'COW',  # class of worker\n",
    "    'MAR',  # marital status\n",
    "    'OCCP', # occupation code\n",
    "    'POBP', # place of birth code\n",
    "    'RELP', # relationship status\n",
    "    'SEX',\n",
    "    'RAC1P', # race code\n",
    "    'DIS',  # disability\n",
    "    'ESP',  # employment status of parents\n",
    "    'CIT',  # citizenship status\n",
    "    'MIG',  # mobility status\n",
    "    'MIL',  # military service\n",
    "    'ANC',  # ancestry\n",
    "    'NATIVITY',\n",
    "    'DEAR',\n",
    "    'DEYE',\n",
    "    'DREM',\n",
    "    'ESR',\n",
    "    'ST',\n",
    "    'FER',\n",
    "    'GCL',\n",
    "    'JWTR',\n",
    "#     'PUMA',\n",
    "#     'POWPUMA',\n",
    "}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "a37a6792",
   "metadata": {},
   "outputs": [],
   "source": [
    "import logging\n",
    "from copy import deepcopy\n",
    "from typing import Tuple\n",
    "from functools import reduce\n",
    "from operator import or_\n",
    "\n",
    "from sklearn.model_selection import train_test_split\n",
    "from folktables import BasicProblem\n",
    "\n",
    "def split_folktables_task(\n",
    "        acs_data: pd.DataFrame,\n",
    "        acs_task: BasicProblem,\n",
    "        train_size: float,\n",
    "        test_size: float,\n",
    "        validation_size: float = None,\n",
    "        max_sensitive_groups: int = None,\n",
    "        stratify_by_state: bool = True,\n",
    "        save_to_disk: Path = None,\n",
    "        file_prefix: str = \"\",\n",
    "        seed: int = 42,\n",
    "    ) -> Tuple[pd.DataFrame, ...]:\n",
    "    \"\"\"Train/test split a given folktables task (for train/test/validation).\n",
    "    \n",
    "    According to the dataset's datasheet, (at least) the ACSIncome\n",
    "    task should be stratified by state.\n",
    "    \n",
    "    Parameters\n",
    "    ----------\n",
    "    acs_data : pd.DataFrame\n",
    "    acs_task : folktables.BasicProblem\n",
    "    train_size : float\n",
    "    test_size : float\n",
    "    validation_size : float\n",
    "    max_sensitive_groups : int, optional\n",
    "        If the number of protected groups exceeds this, discard samples belonging to\n",
    "        the groups with lowest relative size.\n",
    "    stratify_by_state : bool, optional\n",
    "        Whether to stratify splits by state.\n",
    "    seed : int, optional\n",
    "    \n",
    "    Returns\n",
    "    -------\n",
    "    (train_data, test_data, validation_data) : Tuple[pd.DataFrame, pd.DataFrame, pd.DataFrame]\n",
    "    \"\"\"\n",
    "    # Sanity check\n",
    "    assert train_size + test_size + (validation_size or 0.0) == 1\n",
    "    assert all(val is None or 0 <= val <= 1 for val in (train_size, test_size, validation_size))\n",
    "\n",
    "    # Add State to the feature columns so we can do stratified splits (will be removed later)\n",
    "    remove_state_col_later = False # only remove the state column later if we were the ones adding it\n",
    "    if stratify_by_state:\n",
    "        if STATE_COL not in acs_task.features:\n",
    "            acs_task = deepcopy(acs_task) # we're gonna need to change this task object\n",
    "            acs_task.features.append(STATE_COL)\n",
    "            remove_state_col_later = True\n",
    "        else:\n",
    "            remove_state_col_later = False\n",
    "\n",
    "    # Pre-process data + select task-specific features\n",
    "    features, label, group = acs_task.df_to_numpy(acs_data)\n",
    "\n",
    "    # Make a DataFrame with all processed data\n",
    "    df = pd.DataFrame(data=features, columns=acs_task.features)\n",
    "    df[acs_task.target] = label\n",
    "\n",
    "    # Correct column ordering (1st: label, 2nd: group, 3rd and onwards: features)\n",
    "    cols_order = ([acs_task.target, acs_task.group] +\n",
    "        list(set(acs_task.features) - {acs_task.group}))\n",
    "    if remove_state_col_later:\n",
    "        cols_order = [col for col in cols_order if col != STATE_COL]\n",
    "\n",
    "    # Save state_col for stratified split\n",
    "    if stratify_by_state:\n",
    "        state_col_data = df[STATE_COL]\n",
    "\n",
    "    # Enforce correct ordering in df\n",
    "    df = df[cols_order]\n",
    "\n",
    "    # Drop samples from sensitive groups with low relative size\n",
    "    # (e.g., original paper has only White and Black races)\n",
    "    if max_sensitive_groups is not None and max_sensitive_groups > 0:\n",
    "        group_sizes = df.value_counts(acs_task.group, sort=True, ascending=False)\n",
    "        big_groups = group_sizes.index.to_list()[: max_sensitive_groups]\n",
    "\n",
    "        big_groups_filter = reduce(\n",
    "            or_,\n",
    "            [(df[acs_task.group].to_numpy() == g) for g in big_groups],\n",
    "        )\n",
    "        \n",
    "        # Keep only big groups\n",
    "        df = df[big_groups_filter]\n",
    "        state_col_data = state_col_data[big_groups_filter]\n",
    "        \n",
    "        # Group values must be sorted, and start at 0\n",
    "        # (e.g., if we deleted group=2 but kept group=3, the later should now have value 2)\n",
    "        if df[acs_task.group].max() > df[acs_task.group].nunique():\n",
    "            map_to_sequential = {g: idx for g, idx in zip(big_groups, range(len(big_groups)))}\n",
    "            df[acs_task.group] = [map_to_sequential[g] for g in df[acs_task.group]]\n",
    "\n",
    "            logging.warning(f\"Using the following group value mapping: {map_to_sequential}\")\n",
    "            assert df[acs_task.group].max() == df[acs_task.group].nunique() - 1\n",
    "\n",
    "    ## Try to enforce correct types\n",
    "    # All columns should be encoded as integers, dtype=int\n",
    "    types_dict = {\n",
    "        col: int for col in df.columns\n",
    "        if df.dtypes[col] != \"object\"\n",
    "    }\n",
    "    \n",
    "    df = df.astype(types_dict)\n",
    "    # ^ set int types right-away so that categories don't have floating points\n",
    "    \n",
    "    # Set categorical columns to start at value=0! (necessary for sensitive attributes)\n",
    "    for col in (ACS_CATEGORICAL_COLS & set(df.columns)):\n",
    "        df[col] = df[col] - df[col].min()\n",
    "\n",
    "    # Set categorical columns to the correct dtype \"category\"\n",
    "    types_dict.update({\n",
    "        col: \"category\" for col in (ACS_CATEGORICAL_COLS & set(df.columns))\n",
    "        # if df[col].nunique() < 10\n",
    "    })\n",
    "\n",
    "    # Plus the group is definitely categorical\n",
    "    types_dict.update({acs_task.group: \"category\"})\n",
    "    \n",
    "    # And the target is definitely integer\n",
    "    types_dict.update({acs_task.target: int})\n",
    "    \n",
    "    # Set df to correct types\n",
    "    df = df.astype(types_dict)\n",
    "\n",
    "    # ** Split data in train/test/validation **\n",
    "    train_idx, other_idx = train_test_split(\n",
    "        df.index,\n",
    "        train_size=train_size,\n",
    "        stratify=state_col_data if stratify_by_state else None,\n",
    "        random_state=seed,\n",
    "        shuffle=True)\n",
    "\n",
    "    train_df, other_df = df.loc[train_idx], df.loc[other_idx]\n",
    "    assert len(set(train_idx) & set(other_idx)) == 0\n",
    "\n",
    "    # Split validation\n",
    "    if validation_size is not None and validation_size > 0:\n",
    "        new_test_size = test_size / (test_size + validation_size)\n",
    "\n",
    "        val_idx, test_idx = train_test_split(\n",
    "            other_df.index,\n",
    "            test_size=new_test_size,\n",
    "            stratify=state_col_data.loc[other_idx] if stratify_by_state else None,\n",
    "            random_state=seed,\n",
    "            shuffle=True)\n",
    "\n",
    "        val_df, test_df = other_df.loc[val_idx], other_df.loc[test_idx]\n",
    "        assert len(train_idx) + len(val_idx) + len(test_idx) == len(df)\n",
    "        assert np.isclose(len(val_df) / len(df), validation_size)\n",
    "\n",
    "    else:\n",
    "        test_idx = other_idx\n",
    "        test_df = other_df\n",
    "\n",
    "    assert np.isclose(len(train_df) / len(df), train_size)\n",
    "    assert np.isclose(len(test_df) / len(df), test_size)\n",
    "    \n",
    "    # Optionally, save data to disk\n",
    "    # Warning: depends on global notebook variables\n",
    "    if save_to_disk:\n",
    "        subfolder_name = f\"train={train_size:.2}_test={test_size:.2}\"\n",
    "        if validation_size:\n",
    "            subfolder_name = f\"{subfolder_name}_validation={validation_size:.2}\"\n",
    "        if max_sensitive_groups is not None and max_sensitive_groups > 0:\n",
    "            subfolder_name = f\"{subfolder_name}_max-groups={max_sensitive_groups}\"\n",
    "\n",
    "        # Create folder\n",
    "        save_to_disk = save_to_disk / subfolder_name\n",
    "        save_to_disk.mkdir(parents=True, exist_ok=True)\n",
    "\n",
    "        print(f\"Saving data to folder '{str(save_to_disk)}' with prefix '{file_prefix}'.\")\n",
    "        train_df.to_csv(save_to_disk / f\"{file_prefix}.train.csv\", header=True, index_label=\"index\")\n",
    "        test_df.to_csv(save_to_disk / f\"{file_prefix}.test.csv\", header=True, index_label=\"index\")\n",
    "        \n",
    "        if validation_size:\n",
    "            val_df.to_csv(save_to_disk / f\"{file_prefix}.validation.csv\", header=True, index_label=\"index\")\n",
    "\n",
    "    return (train_df, test_df, val_df) if validation_size else (train_df, test_df)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6d0b1d1b",
   "metadata": {},
   "outputs": [],
   "source": [
    "import folktables\n",
    "from tqdm.auto import tqdm\n",
    "\n",
    "all_acs_tasks = [\n",
    "    'ACSIncome',\n",
    "    'ACSPublicCoverage',\n",
    "    'ACSMobility',\n",
    "    'ACSEmployment',\n",
    "    'ACSTravelTime',\n",
    "]\n",
    "\n",
    "const_predictor_acc = {}\n",
    "\n",
    "# Generate data and save to disk, for all tasks\n",
    "for task_name in tqdm(all_acs_tasks):\n",
    "\n",
    "    # Dynamically import/load task object\n",
    "    task_obj = getattr(folktables, task_name)\n",
    "\n",
    "    # Process data\n",
    "    data = split_folktables_task(\n",
    "        acs_data,\n",
    "        task_obj,\n",
    "        train_size=TRAIN_SIZE,\n",
    "        test_size=TEST_SIZE,\n",
    "        validation_size=VALIDATION_SIZE,\n",
    "        max_sensitive_groups=MAX_SENSITIVE_GROUPS,\n",
    "        stratify_by_state=True,\n",
    "        seed=SEED,\n",
    "        save_to_disk=data_dir,\n",
    "        file_prefix=task_name,\n",
    "    )\n",
    "    \n",
    "    const_predictor_acc[task_name] = {\n",
    "        curr_type: max(curr_data[task_obj.target].mean(), 1-curr_data[task_obj.target].mean())\n",
    "        for curr_type, curr_data in zip([\"train\", \"test\", \"validation\"], data)\n",
    "    }"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1c550577",
   "metadata": {},
   "source": [
    "## Log the constant classifier accuracy for each dataset and data type\n",
    "The constant classifier always predicts either class 1 or 0 (whichever has highest prevalence in the dataset)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "5bc9927c",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{\n",
      "  \"ACSIncome\": {\n",
      "    \"train\": 0.6276484057384071,\n",
      "    \"test\": 0.6267948098880045,\n",
      "    \"validation\": 0.6275487157120827\n",
      "  },\n",
      "  \"ACSPublicCoverage\": {\n",
      "    \"train\": 0.705777716173447,\n",
      "    \"test\": 0.7080615054796799,\n",
      "    \"validation\": 0.7048998544597685\n",
      "  },\n",
      "  \"ACSMobility\": {\n",
      "    \"train\": 0.7366843896876353,\n",
      "    \"test\": 0.7369037614281713,\n",
      "    \"validation\": 0.7358002054311933\n",
      "  },\n",
      "  \"ACSEmployment\": {\n",
      "    \"train\": 0.5401004328892394,\n",
      "    \"test\": 0.5415303116248027,\n",
      "    \"validation\": 0.5397890715398488\n",
      "  },\n",
      "  \"ACSTravelTime\": {\n",
      "    \"train\": 0.5622398192707292,\n",
      "    \"test\": 0.5624768352647576,\n",
      "    \"validation\": 0.5620214827548473\n",
      "  }\n",
      "}\n"
     ]
    }
   ],
   "source": [
    "import json\n",
    "print(json.dumps(const_predictor_acc, indent=2))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "714cd7d7",
   "metadata": {},
   "source": [
    "---"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.10.12"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}