00_randomize_train_test.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Train/test random split\n",
    "\n",
    "In this notebook, we define the nights that will  be used in the training / testing set and we homogenize demographics data across all datasets."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os, glob\n",
    "import numpy as np\n",
    "import pandas as pd\n",
    "\n",
    "# Define root data paths, where the NSRR data is stored\n",
    "root_dir = '/Volumes/NSRR/'\n",
    "assert os.path.isdir(root_dir)\n",
    "\n",
    "# Define the output path\n",
    "out_dir = 'output/demo/'\n",
    "if not os.path.isdir(out_dir):\n",
    "    os.mkdir(out_dir)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### CCSHS\n",
    "\n",
    "https://sleepdata.org/datasets/ccshs"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Load demographics CSV file\n",
    "df_ccshs = pd.read_csv(\n",
    "    root_dir + 'ccshs/datasets/ccshs-trec-dataset-0.6.0.csv', \n",
    "    usecols=['nsrrid', 'age', 'ethnicity', 'male', 'bmi', 'race3', 'ahi_a0h3', 'yinsdx', 'yhtn', 'ydiab', 'ydep']\n",
    ")\n",
    "\n",
    "# Rename \n",
    "df_ccshs['race3'].replace({1: 'caucasian', 2: 'african', 3: 'other'}, \n",
    "                          inplace=True)\n",
    "df_ccshs.loc[df_ccshs['ethnicity'] == 1, 'race3'] = 'hispanic'\n",
    "df_ccshs.drop(columns=['ethnicity'], inplace=True)\n",
    "df_ccshs.rename(columns={\n",
    "    'race3': 'ethnicity', \n",
    "    'nsrrid': 'subj', \n",
    "    'ahi_a0h3': 'ahi',\n",
    "    'yinsdx': 'insomnia',\n",
    "    'yhtn': 'hypertension',\n",
    "    'ydiab': 'diabete',\n",
    "    'ydep': 'depression'}, inplace=True)\n",
    "\n",
    "# Remove subject for which we don't have the data\n",
    "df_ccshs = df_ccshs[~df_ccshs['subj'].isin([1800639, 1800842])]\n",
    "\n",
    "# Convert to str\n",
    "df_ccshs['subj'] = df_ccshs['subj'].astype(str)\n",
    "df_ccshs.set_index('subj', inplace=True)\n",
    "\n",
    "# Define training / testing\n",
    "# Keep only a random subset of 415 subjects for training (leave 100 for testing)\n",
    "df_ccshs[\"set\"] = \"excluded\"\n",
    "idx_train = df_ccshs.sample(n=415, replace=False, random_state=42).index\n",
    "idx_test = np.setdiff1d(df_ccshs.index, idx_train)\n",
    "df_ccshs.loc[idx_train, \"set\"] = \"training\"\n",
    "df_ccshs.loc[idx_test, \"set\"] = \"testing\"\n",
    "\n",
    "# Export demographics to CSV file\n",
    "df_ccshs['dataset'] = 'CCSHS'\n",
    "# df_ccshs.to_csv(out_dir + \"demo_nsrr_ccshs.csv\")\n",
    "\n",
    "print(df_ccshs.shape[0], \"unique nights.\")\n",
    "print(df_ccshs['set'].value_counts())\n",
    "df_ccshs.head()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### CFS\n",
    "\n",
    "https://sleepdata.org/datasets/cfs"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "df_cfs = pd.read_csv(\n",
    "    root_dir + 'cfs/datasets/cfs-visit5-dataset-0.4.0.csv', \n",
    "    usecols=['nsrrid', 'age', 'SEX', 'race', 'ethnicity', 'bmi', 'ahi_a0h3',\n",
    "             'AbnorEEG', 'AbnorEye', 'QuEEG1', 'QuHR', 'QuEOGL', 'QuChin', \n",
    "             'RemNRemPr', 'Stg1Stg2Pr', 'Stg2Stg3Pr', 'WakSlePr', \n",
    "             'INSODIAG', 'htndx', 'DIADIAG', 'DIANARC', 'DEPDIAG'])\n",
    "\n",
    "# Rename for ease-of-use\n",
    "df_cfs.rename(columns={\n",
    "    'nsrrid': 'subj',\n",
    "    'SEX': 'male',\n",
    "    'ahi_a0h3': 'ahi',\n",
    "    'INSODIAG': 'insomnia',\n",
    "    'htndx': 'hypertension',\n",
    "    'DIADIAG': 'diabete',\n",
    "    'DIANARC': 'narcolepsy',\n",
    "    'DEPDIAG': 'depression'\n",
    "}, inplace=True)\n",
    "\n",
    "df_cfs['race'].replace({1: 'caucasian', 2: 'african', 3: 'other'}, inplace=True)\n",
    "df_cfs.loc[df_cfs['ethnicity'] == 1, 'race'] = 'hispanic'\n",
    "df_cfs.drop(columns=['ethnicity'], inplace=True)\n",
    "df_cfs.rename(columns={'race': 'ethnicity'}, inplace=True)\n",
    "\n",
    "# Remove subjects with no EEG / hypnogram file\n",
    "df_cfs = df_cfs[~df_cfs['subj'].isin([800269, 802230, 801601, 801811, 801841, 801953, 802234])]\n",
    "\n",
    "# Remove Ss with bad data quality or unreliable scoring (DISABLED)\n",
    "# df_cfs = df_cfs[df_cfs['RemNRemPr'] != 1]\n",
    "# df_cfs = df_cfs[df_cfs['Stg1Stg2Pr'] != 1]\n",
    "# df_cfs = df_cfs[df_cfs['Stg2Stg3Pr'] != 1]\n",
    "# df_cfs = df_cfs[df_cfs['WakSlePr'] != 1]\n",
    "# df_cfs = df_cfs[df_cfs['AbnorEEG'] != 1]\n",
    "# df_cfs = df_cfs[df_cfs['AbnorEye'] != 1]\n",
    "# df_cfs = df_cfs[df_cfs['QuEEG1'] >= 4] # >= 4 = signal good for > 75% of the time\n",
    "# df_cfs = df_cfs[df_cfs['QuEOGL'] >= 4]\n",
    "# df_cfs = df_cfs[df_cfs['QuChin'] >= 4]\n",
    "\n",
    "# Drop columns\n",
    "df_cfs.drop(\n",
    "    columns=['RemNRemPr', 'AbnorEEG', 'AbnorEye', 'QuEEG1', 'QuHR', 'QuEOGL', \n",
    "             'QuChin', 'RemNRemPr', 'Stg1Stg2Pr', 'Stg2Stg3Pr', 'WakSlePr'], \n",
    "    inplace=True)\n",
    "\n",
    "# Convert to str\n",
    "df_cfs['subj'] = df_cfs['subj'].astype(str)\n",
    "df_cfs.set_index('subj', inplace=True)\n",
    "\n",
    "# Define training / testing\n",
    "# Keep only a random subset of 600 subjects for training to avoid dataset imbalance\n",
    "df_cfs[\"set\"] = \"excluded\"\n",
    "idx_train = df_cfs.sample(n=600, replace=False, random_state=42).index\n",
    "idx_test = np.setdiff1d(df_cfs.index, idx_train)\n",
    "# Now we keep 100 random participants of ``idx_test`` for testing\n",
    "rs = np.random.RandomState(42)\n",
    "idx_test = rs.choice(idx_test, size=100, replace=False)\n",
    "df_cfs.loc[idx_train, \"set\"] = \"training\"\n",
    "df_cfs.loc[idx_test, \"set\"] = \"testing\"\n",
    "\n",
    "# Export demographics to CSV file\n",
    "df_cfs['dataset'] = 'CFS'\n",
    "# df_cfs.to_csv(out_dir + \"demo_nsrr_cfs.csv\")\n",
    "\n",
    "print(df_cfs.shape[0], \"unique nights.\")\n",
    "print(df_cfs['set'].value_counts())\n",
    "df_cfs.head()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### CHAT\n",
    "\n",
    "https://sleepdata.org/datasets/chat"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "usecols = ['nsrrid', 'male', 'ageyear_at_meas', 'overall', \n",
    "           'ant5', 'chi3', 'race3', 'ahi_a0h3', 'med1o2', 'med1m2', 'med1i2', 'med1h2']\n",
    "\n",
    "df_chat = pd.read_csv(\n",
    "    root_dir + 'chat/datasets/chat-baseline-dataset-0.11.0.csv', usecols=usecols)\n",
    "\n",
    "# Rename columns\n",
    "df_chat.rename(columns={'nsrrid': 'subj',\n",
    "                        'ageyear_at_meas': 'age',\n",
    "                        'ant5': 'bmi',\n",
    "                        'chi3': 'ethnicity',\n",
    "                        'race3': 'race',\n",
    "                        'ahi_a0h3': 'ahi',\n",
    "                        'med1o2': 'insomnia',\n",
    "                        'med1m2': 'hypertension',\n",
    "                        'med1i2': 'diabete',\n",
    "                        'med1h2': 'depression'\n",
    "                      }, inplace=True)\n",
    "\n",
    "df_chat['race'].replace({1: 'caucasian', 2: 'african', 3: 'other'}, inplace=True)\n",
    "df_chat.loc[df_chat['ethnicity'] == 1, 'race'] = 'hispanic'\n",
    "df_chat.drop(columns=['ethnicity'], inplace=True)\n",
    "df_chat.rename(columns={'race': 'ethnicity'}, inplace=True)\n",
    "\n",
    "# Keep only \"Excellent\" quality study (DISABLED)\n",
    "# print(df_chat[df_chat['overall'] < 6].shape[0], 'subjects with bad PSG data quality will be removed.')\n",
    "# df_chat = df_chat[df_chat['overall'] >= 6]\n",
    "\n",
    "# Remove nights for which sex is NaN\n",
    "df_chat.dropna(subset=['male'], inplace=True)\n",
    "\n",
    "# Convert to str\n",
    "df_chat['subj'] = df_chat['subj'].apply(lambda x: str(x).zfill(4))\n",
    "df_chat.set_index('subj', inplace=True)\n",
    "df_chat['male'] = df_chat['male'].astype(int)\n",
    "df_chat['age'] = df_chat['age'].astype(float)\n",
    "\n",
    "# Define training / testing\n",
    "# Keep only a random subset of ~350 subjects for training (leave 100 for testing)\n",
    "df_chat[\"set\"] = \"excluded\"\n",
    "idx_train = df_chat.sample(n=353, replace=False, random_state=42).index\n",
    "idx_test = np.setdiff1d(df_chat.index, idx_train)\n",
    "# Now we keep 100 random participants of ``idx_test`` for testing\n",
    "rs = np.random.RandomState(42)\n",
    "idx_test = rs.choice(idx_test, size=100, replace=False)\n",
    "df_chat.loc[idx_train, \"set\"] = \"training\"\n",
    "df_chat.loc[idx_test, \"set\"] = \"testing\"\n",
    "\n",
    "# Export demographics to CSV file\n",
    "df_chat['dataset'] = 'CHAT'\n",
    "# df_chat.to_csv(out_dir + \"demo_nsrr_chat.csv\")\n",
    "\n",
    "print(df_chat.shape[0], \"unique nights.\")\n",
    "print(df_chat['set'].value_counts())\n",
    "df_chat.head()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### HomePAP\n",
    "\n",
    "https://sleepdata.org/datasets/homepap"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "desc_dir = root_dir + 'homepap/datasets/homepap-baseline-dataset-0.1.0.csv'\n",
    "hypno_dir = root_dir + 'homepap/polysomnography/annotations-events-profusion/lab/full/'\n",
    "usecols = ['nsrrid', 'treatmentarm', 'age', 'gender', 'bmi', 'ahi_full', 'race3', 'ethnicity',\n",
    "           'dxhtn', 'dxdiab', 'dxdep']\n",
    "\n",
    "df_hpap = pd.read_csv(desc_dir, usecols=usecols)\n",
    "\n",
    "df_hpap.rename(columns={\n",
    "    'nsrrid': 'subj',\n",
    "    'race3': 'race',\n",
    "    'ahi_full': 'ahi',\n",
    "    'dxhtn': 'hypertension',\n",
    "    'dxdiab': 'diabete',\n",
    "    'dxdep': 'depression'\n",
    "    }, inplace=True)\n",
    "\n",
    "df_hpap['race'].replace({1: 'caucasian', 2: 'african', 3: 'other'}, inplace=True)\n",
    "df_hpap.loc[df_hpap['ethnicity'] == 1, 'race'] = 'hispanic'\n",
    "df_hpap.drop(columns=['ethnicity'], inplace=True)\n",
    "df_hpap.rename(columns={'race': 'ethnicity'}, inplace=True)\n",
    "\n",
    "df_hpap['male'] = df_hpap['gender'].replace({'M': 1, 'F': 0})\n",
    "\n",
    "# Drop NaN\n",
    "# df_hpap.dropna(inplace=True)\n",
    "\n",
    "# Keep only in-lab PSG\n",
    "df_hpap = df_hpap[df_hpap['treatmentarm'] == 1]\n",
    "df_hpap.drop(columns=['treatmentarm', 'gender'], inplace=True)\n",
    "\n",
    "# Keep only full night\n",
    "files = [int(f.split('-')[3]) for f in os.listdir(hypno_dir)]\n",
    "good = np.intersect1d(files, df_hpap['subj'].values, assume_unique=True)\n",
    "df_hpap = df_hpap[df_hpap['subj'].isin(good)]\n",
    "\n",
    "# Convert to str\n",
    "df_hpap['subj'] = df_hpap['subj'].astype(str)\n",
    "df_hpap.set_index('subj', inplace=True)\n",
    "\n",
    "# Because of the few number of participants, we only use HomePaP as a training dataset\n",
    "df_hpap[\"set\"] = \"training\"\n",
    "\n",
    "# Export demographics to CSV file\n",
    "df_hpap['dataset'] = 'HOMEPAP'\n",
    "# df_hpap.to_csv(out_dir + \"demo_nsrr_homepap.csv\")\n",
    "\n",
    "print(df_hpap.shape[0], 'subjects remaining')\n",
    "print(df_hpap['set'].value_counts())\n",
    "df_hpap.head(10)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### MESA\n",
    "\n",
    "https://sleepdata.org/datasets/mesa"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# BioLinCC\n",
    "df_mesa_bio = pd.read_csv(\"data/mesae5_drepos_20151101.csv\", usecols=['mesaid', 'bmi5c', 'htn5c', 'dm035c'])\n",
    "\n",
    "# Diabetes: 0=normal, 1=impaired, 2=treated diabete, 3=untreated diabete\n",
    "df_mesa_bio.rename(\n",
    "    columns={'mesaid': 'subj', 'bmi5c': 'bmi', 'htn5c': 'hypertension', 'dm035c': 'diabete'}, inplace=True)\n",
    "\n",
    "df_mesa_bio"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "desc_dir = root_dir + 'mesa/datasets/mesa-sleep-dataset-0.3.0.csv'\n",
    "\n",
    "# No BMI for MESA -- see BioLinCC instead\n",
    "df_mesa = pd.read_csv(\n",
    "    desc_dir, usecols=['mesaid', 'gender1', 'sleepage5c', 'overall5', 'race1c', 'ahi_a0h3', 'insmnia5'])\n",
    "\n",
    "# Rename columns\n",
    "df_mesa.rename(columns={'mesaid': 'subj',\n",
    "                        'gender1': 'gender',\n",
    "                        'sleepage5c': 'age',\n",
    "                        'overall5': 'overall',\n",
    "                        'race1c': 'ethnicity',\n",
    "                        'ahi_a0h3': 'ahi',\n",
    "                        'insmnia5': 'insomnia',\n",
    "                      }, inplace=True)\n",
    "\n",
    "df_mesa['ethnicity'].replace({1: 'caucasian', 2: 'other', 3: 'african', 4: 'hispanic'}, inplace=True)\n",
    "\n",
    "# Keep only \"Excellent\" quality study (DISABLED)\n",
    "# print(df_mesa[df_mesa['overall'] < 6].shape[0], \n",
    "#       'subjects with bad PSG data quality will be removed.')\n",
    "# df_mesa = df_mesa[df_mesa['overall'] >= 6]\n",
    "df_mesa['male'] = (df_mesa['gender'] == 1).astype(int)\n",
    "df_mesa.drop(columns=['gender'], inplace=True)\n",
    "\n",
    "# Merge with BioLinCC\n",
    "df_mesa = df_mesa.merge(df_mesa_bio, how=\"left\")\n",
    "\n",
    "# Convert to str\n",
    "df_mesa['subj'] = df_mesa['subj'].apply(lambda x: str(x).zfill(4))\n",
    "df_mesa.set_index('subj', inplace=True)\n",
    "\n",
    "# In MESA, there are a large number of EDF files that are missing\n",
    "# We therefore only keep subjects that have valid data here.\n",
    "edf_files = sorted(glob.glob(root_dir + 'mesa/polysomnography/edfs/*.edf'))\n",
    "edf_files = [os.path.basename(c)[-8:-4] for c in edf_files]\n",
    "df_mesa = df_mesa.loc[edf_files] \n",
    "\n",
    "# Define training / testing\n",
    "# Keep only a random subset of 600 subjects for training to avoid dataset imbalance\n",
    "df_mesa[\"set\"] = \"excluded\"\n",
    "idx_train = df_mesa.sample(n=600, replace=False, random_state=42).index\n",
    "idx_test = np.setdiff1d(df_mesa.index, idx_train)\n",
    "# Now we keep 100 random participants of ``idx_test`` for testing\n",
    "rs = np.random.RandomState(42)\n",
    "idx_test = rs.choice(idx_test, size=100, replace=False)\n",
    "df_mesa.loc[idx_train, \"set\"] = \"training\"\n",
    "df_mesa.loc[idx_test, \"set\"] = \"testing\"\n",
    "\n",
    "# Export demographics to CSV file\n",
    "df_mesa['dataset'] = 'MESA'\n",
    "# df_mesa.to_csv(out_dir + \"demo_nsrr_mesa.csv\")\n",
    "\n",
    "print(df_mesa.shape[0], 'subjects remaining')\n",
    "print(df_mesa['set'].value_counts())\n",
    "df_mesa.head(10)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### MrOS\n",
    "\n",
    "https://sleepdata.org/datasets/mros"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "desc_dir = root_dir + 'mros/datasets/mros-visit1-dataset-0.4.0.csv'\n",
    "usecols = ['nsrrid', 'visit', 'gender', 'vsage1', 'pooveral', 'hwbmi', 'gierace', 'poordi3', \n",
    "           'slinsomn', 'mhbp', 'mhdiab', 'slnarc']\n",
    "\n",
    "df_mros = pd.read_csv(desc_dir, usecols=usecols, low_memory=False)\n",
    "\n",
    "# Rename columns\n",
    "df_mros.rename(columns={'nsrrid': 'subj',\n",
    "                        'vsage1': 'age',\n",
    "                        'pooveral': 'overall',\n",
    "                        'hwbmi': 'bmi',\n",
    "                        'gierace': 'ethnicity',\n",
    "                        'poordi3': 'ahi',\n",
    "                        'slinsomn': 'insomnia',\n",
    "                        'mhbp': 'hypertension',\n",
    "                        'mhdiab': 'diabete',\n",
    "                        'slnarc': 'narcolepsy',\n",
    "                      }, inplace=True)\n",
    "\n",
    "# Keep only \"Excellent\" quality study\n",
    "# print(df_mros[df_mros['overall'] < 6].shape[0], \n",
    "#       'subjects with bad PSG data quality will be removed.')\n",
    "# df_mros = df_mros[df_mros['overall'] >= 6]\n",
    "\n",
    "# They should all be male in MrOS!\n",
    "df_mros['male'] = (df_mros['gender'] == 2).astype(int)\n",
    "df_mros['ethnicity'].replace({1: 'caucasian', 2: 'african', 3: 'other', \n",
    "                              4: 'hispanic', 5: 'other'}, inplace=True)\n",
    "\n",
    "# Keep only first visit\n",
    "df_mros = df_mros[df_mros['visit'] == 1]\n",
    "\n",
    "# Convert to str\n",
    "df_mros['subj'] = df_mros['subj'].apply(lambda x: str(x).zfill(4))\n",
    "df_mros.set_index('subj', inplace=True)\n",
    "\n",
    "# Fix an invalid 'A' in BMI\n",
    "df_mros.loc[df_mros['bmi'] == 'A', 'bmi'] = np.nan\n",
    "df_mros['bmi'] = df_mros['bmi'].astype(float)\n",
    "\n",
    "# Define training / testing\n",
    "# Keep only a random subset of 600 subjects for training to avoid dataset imbalance\n",
    "df_mros[\"set\"] = \"excluded\"\n",
    "idx_train = df_mros.sample(n=600, replace=False, random_state=42).index\n",
    "idx_test = np.setdiff1d(df_mros.index, idx_train)\n",
    "# Now we keep 100 random participants of ``idx_test`` for testing\n",
    "rs = np.random.RandomState(42)\n",
    "idx_test = rs.choice(idx_test, size=100, replace=False)\n",
    "df_mros.loc[idx_train, \"set\"] = \"training\"\n",
    "df_mros.loc[idx_test, \"set\"] = \"testing\"\n",
    "\n",
    "# Export demographics to CSV file\n",
    "df_mros['dataset'] = 'MROS'\n",
    "# df_mros.to_csv(out_dir + \"demo_nsrr_mros.csv\")\n",
    "\n",
    "print(df_mros.shape[0], 'subjects remaining')\n",
    "print(df_mros['set'].value_counts())\n",
    "df_mros.head(10)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### SHHS\n",
    "\n",
    "https://sleepdata.org/datasets/shhs"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "desc_dir = root_dir + 'shhs/datasets/shhs1-dataset-0.16.0.csv'\n",
    "usecols = ['nsrrid', 'visitnumber', 'gender', 'age_s1', 'overall_shhs1', \n",
    "           'race', 'bmi_s1', 'ahi_a0h3', 'ethnicity', 'HTNDerv_s1', 'ParRptDiab']\n",
    "\n",
    "df_shhs = pd.read_csv(desc_dir, usecols=usecols)\n",
    "\n",
    "# Rename columns\n",
    "df_shhs.rename(columns={'nsrrid': 'subj',\n",
    "                        'gender1': 'gender',\n",
    "                        'age_s1': 'age',\n",
    "                        'overall_shhs1': 'overall',\n",
    "                        'bmi_s1': 'bmi',\n",
    "                        'ahi_a0h3': 'ahi',\n",
    "                        'HTNDerv_s1': 'hypertension',\n",
    "                        'ParRptDiab': 'diabete'\n",
    "                      }, inplace=True)\n",
    "\n",
    "df_shhs['race'].replace({1: 'caucasian', 2: 'african', 3: 'other'}, inplace=True)\n",
    "df_shhs.loc[df_shhs['ethnicity'] == 1, 'race'] = 'hispanic'\n",
    "df_shhs.drop(columns=['ethnicity'], inplace=True)\n",
    "df_shhs.rename(columns={'race': 'ethnicity'}, inplace=True)\n",
    "\n",
    "# Keep only \"Excellent\" quality study\n",
    "# print(df_shhs[df_shhs['overall'] < 6].shape[0], \n",
    "#       'subjects with bad PSG data quality will be removed.')\n",
    "# df_shhs = df_shhs[df_shhs['overall'] >= 6]\n",
    "\n",
    "df_shhs['male'] = (df_shhs['gender'] == 1).astype(int)\n",
    "\n",
    "# Keep only first visit\n",
    "df_shhs = df_shhs[df_shhs['visitnumber'] == 1]\n",
    "\n",
    "# Convert to str\n",
    "df_shhs['subj'] = df_shhs['subj'].apply(lambda x: str(x).zfill(4))\n",
    "df_shhs.set_index('subj', inplace=True)\n",
    "\n",
    "# Define training / testing\n",
    "# Keep only a random subset of 600 subjects for training to avoid dataset imbalance\n",
    "df_shhs[\"set\"] = \"excluded\"\n",
    "idx_train = df_shhs.sample(n=600, replace=False, random_state=42).index\n",
    "idx_test = np.setdiff1d(df_shhs.index, idx_train)\n",
    "# Now we keep 100 random participants of ``idx_test`` for testing\n",
    "rs = np.random.RandomState(42)\n",
    "idx_test = rs.choice(idx_test, size=100, replace=False)\n",
    "df_shhs.loc[idx_train, \"set\"] = \"training\"\n",
    "df_shhs.loc[idx_test, \"set\"] = \"testing\"\n",
    "\n",
    "# Export demographics to CSV file\n",
    "df_shhs['dataset'] = 'SHHS'\n",
    "# df_shhs.to_csv(out_dir + \"demo_nsrr_shhs.csv\")\n",
    "\n",
    "print(df_shhs.shape[0], 'subjects remaining')\n",
    "print(df_shhs['set'].value_counts())\n",
    "df_shhs.head(10)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Combine and export"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "df_all = pd.concat([df_ccshs, df_cfs, df_chat, df_hpap, df_mesa, df_mros, df_shhs])\n",
    "df_all.drop(columns=['gender', 'visit', 'visitnumber', 'overall'], inplace=True)\n",
    "df_all = df_all.set_index(\"dataset\", append=True).reorder_levels([\"dataset\", \"subj\"])\n",
    "\n",
    "# Remove \"excluded\"\n",
    "df_all = df_all[df_all[\"set\"] != \"excluded\"]\n",
    "\n",
    "df_all"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**Column specific preprocessing**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Convert age to int\n",
    "df_all['age'] = df_all['age'].round().astype(int)\n",
    "df_all['age'].hist(bins=np.arange(5, 100, 5));"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Insomnia\n",
    "df_all['insomnia'] = df_all['insomnia'].replace({\"M\": np.nan, \"K\": np.nan, -2: np.nan}).astype(float)\n",
    "df_all['insomnia'].value_counts(dropna=False)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Hypertension\n",
    "df_all['hypertension'] = df_all['hypertension'].astype(float)\n",
    "df_all['hypertension'].value_counts(dropna=False)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Diabete\n",
    "df_all['diabete'] = df_all['diabete'].replace({2: 1, 3: 1, -2: np.nan}).astype(float)\n",
    "df_all['diabete'].value_counts(dropna=False)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Narcolepsy\n",
    "df_all['narcolepsy'] = df_all['narcolepsy'].replace({\"M\": np.nan, \"K\": np.nan, -2: np.nan}).astype(float)\n",
    "df_all['narcolepsy'].value_counts(dropna=False)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Depression\n",
    "df_all['depression'] = df_all['depression'].replace({-2: np.nan}).astype(float)\n",
    "df_all['depression'].value_counts(dropna=False)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**Export**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Re-order columns\n",
    "cols_order = [\n",
    "    'age', 'male', 'bmi', 'ahi', 'ethnicity', 'set', 'depression', 'diabete', \n",
    "    'hypertension', 'insomnia', 'narcolepsy']\n",
    "df_all = df_all[cols_order]\n",
    "df_all.head().round(2)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Export to .csv\n",
    "df_all.to_csv(out_dir + \"demo_nsrr_all.csv\", index=True)"
   ]
  }
 ],
 "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.9.6"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}