HRS_exploration.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "%matplotlib inline\n",
    "import matplotlib.pyplot as plt\n",
    "import collections\n",
    "import gzip\n",
    "import numpy as np\n",
    "import os\n",
    "import phate\n",
    "import time\n",
    "\n",
    "import pandas as pd\n",
    "\n",
    "from collections import defaultdict\n",
    "import itertools\n",
    "\n",
    "# Interactive HTML tools\n",
    "from ipywidgets import interact\n",
    "import bokeh\n",
    "import bokeh.io\n",
    "from bokeh.io import push_notebook\n",
    "from bokeh.plotting import figure, show, save, output_notebook, output_file\n",
    "from bokeh.palettes import Category20b\n",
    "from bokeh.palettes import Category10\n",
    "\n",
    "# Machine-learning and dimensionality reduction tools\n",
    "import sklearn\n",
    "from sklearn import decomposition\n",
    "from sklearn.decomposition import PCA as PCA # We'll use this to check our implementation\n",
    "from sklearn.manifold import TSNE\n",
    "from sklearn.decomposition import TruncatedSVD\n",
    "from sklearn.manifold import MDS"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Notes\n",
    "\n",
    "This notebook is intended for HRS data. It assumes the relevant PCs have already been created (I did this in PLINK).\n",
    "\n",
    "QC reports on genotype data available here: https://hrs.isr.umich.edu/data-products/genetic-data/products#gdv1"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Define functions"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "# Define a function that will allow us to project using TSNE, given inputs and perplexities, and output the coordinates\n",
    "def hrs_pc_tsne(in_data, plex, subset_group, save=False, subsetting=False):\n",
    "    num_pcs = in_data.shape[1]\n",
    "    temp_proj = TSNE(n_components=2, perplexity=plex).fit_transform(in_data)\n",
    "    \n",
    "    if save == True:\n",
    "        if subsetting == True:\n",
    "            np.savetxt('HRS_TSNE_' + subset_group + 'PC_' + str(num_pcs) + '_PLEX_' + str(plex), temp_proj)\n",
    "        elif subsetting == False:\n",
    "            np.savetxt('HRS_TSNE_PC_' + str(num_pcs) + '_PLEX_' + str(plex), temp_proj)\n",
    "    \n",
    "    return(temp_proj)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Importing data\n",
    "\n",
    "Import the principal components and any related auxiliary data."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "# Note: PLINK command used to create cleaned HRS data \n",
    "# plink --bfile HRS --maf 0.05 --mind 0.1 --geno 0.1 --hwe 1e-6 --make-bed --out HRS_CLEAR\n",
    "# Minor Allele Freq, Missing genotype rates (per-sample) Missing genotype rates (per-variant), \n",
    "# Hardy-Weinberg equilibrium, \n",
    "\n",
    "# PLINK command used to create PCs\n",
    "# plink --bfile HRS_CLEAR --pca [n]\n",
    "\n",
    "# PLINK command used to create PCs for white-identified individuals in HRS\n",
    "# Other subsets (black, hispanic) were done in the same manner\n",
    "# plink --bfile HRS_CLEAR --pca [n] --keep hrs_subset_white.txt\n",
    "\n",
    "hrs_data_dir = '/Volumes/Stockage/alex/hrs/projections'\n",
    "\n",
    "# Define the files we'll be using\n",
    "aux_file = 'allIndivs_filtered.txt'\n",
    "pc_files = ['plink.eigenvec_200','HRS_PCA_black.eigenvec','HRS_PCA_hispanic.eigenvec','HRS_PCA_white.eigenvec']\n",
    "pc_file = 'plink.eigenvec_200'\n",
    "\n",
    "aux_path = os.path.join(hrs_data_dir, aux_file)\n",
    "pc_path = os.path.join(hrs_data_dir, pc_file)\n",
    "\n",
    "# Import auxiliary data. Contains IDs and demographic information.\n",
    "# NOTE: The auxiliary data is sorted in an order different from the PC data.\n",
    "aux_data = []\n",
    "with open(aux_path) as input_file:\n",
    "    for line in input_file:\n",
    "        aux_data.append(line.strip().split(','))\n",
    "\n",
    "# Import PC data. This data must be converted to an array.\n",
    "with open(pc_path) as h:\n",
    "    hrs_contents = h.readlines()\n",
    "\n",
    "hrs_data = []\n",
    "\n",
    "for h in hrs_contents:\n",
    "    hrs_data.append(h.split()[2:len(h)])\n",
    "\n",
    "hrs_data_array = np.array(hrs_data).astype(np.float)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "# This is a more general version of above using various subsets of HRS data\n",
    "hrs_data_dir = '/Volumes/Stockage/alex/hrs/projections'\n",
    "\n",
    "# Define the files we'll be using\n",
    "aux_file = 'allIndivs_filtered.txt'\n",
    "\n",
    "# Set up a dict containing subsets and their related files.\n",
    "hrs_subsets_dict = {'All':'plink.eigenvec_200',\n",
    "                    'Black':'HRS_PCA_black.eigenvec',\n",
    "                   'Hispanic':'HRS_PCA_hispanic.eigenvec',\n",
    "                    'White':'HRS_PCA_white.eigenvec'}\n",
    "\n",
    "aux_path = os.path.join(hrs_data_dir, aux_file)\n",
    "\n",
    "aux_data = []\n",
    "with open(aux_path) as input_file:\n",
    "    for line in input_file:\n",
    "        aux_data.append(line.strip().split(','))\n",
    "\n",
    "hrs_data_dict = defaultdict(np.array)\n",
    "\n",
    "for p in hrs_subsets_dict:\n",
    "    pc_path = os.path.join(hrs_data_dir, hrs_subsets_dict[p])\n",
    "    \n",
    "    # Import PC data. This data must be converted to an array.\n",
    "    with open(pc_path) as h:\n",
    "        hrs_contents = h.readlines()\n",
    "\n",
    "    hrs_data = []\n",
    "        \n",
    "    for h in hrs_contents:\n",
    "        hrs_data.append(h.split()[2:len(h)])\n",
    "\n",
    "    hrs_data_dict[p] = np.array(hrs_data).astype(np.float)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "for p in hrs_subsets_dict:\n",
    "    print(p)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "hrs_data_array.shape"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Import ADMIXTURE work"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "admix_dir = '/Volumes/Stockage/alex/hrs/Admixture'\n",
    "\n",
    "# Separated into two datasets (roughly speaking, African Americans and European Americans)\n",
    "admix_data = []\n",
    "\n",
    "# Import data from first file\n",
    "with open(os.path.join(admix_dir,'HRS.txt')) as input_file:\n",
    "    for line in input_file:\n",
    "        admix_data.append(line.strip().split())\n",
    "\n",
    "# Import data from second file\n",
    "with open(os.path.join(admix_dir,'HRS_EurAm.txt')) as input_file:\n",
    "    for line in input_file:\n",
    "        admix_data.append(line.strip().split())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "# Transform this into a Pandas data frame and convert the numeric values to... numeric\n",
    "# ID is stored as numeric as well. This is to match with HRS data, which is sorted numerically by IDs\n",
    "admix_df = pd.DataFrame.from_records(admix_data, columns=['ID','ADMIX1','ADMIX2','ADMIX3'])\n",
    "admix_df[['ADMIX1','ADMIX2','ADMIX3']] = admix_df[['ADMIX1','ADMIX2','ADMIX3']].apply(pd.to_numeric)\n",
    "#admix_df.ID = admix_df.ID.astype(np.int64)\n",
    "admix_df['ID'] = admix_df.ID.astype(np.int64)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Prep work\n",
    "\n",
    "The data is not consistently sorted between files. Also we need to create a variety of dictionaries and their reversals so we can correctly label points later. Also this should probably be called \"Labelling work\" or something."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "# Select what subset to use for analysis:\n",
    "subset = 'All'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "# Import the data and convert it to a pandas dataset for easier multivariable sorting.\n",
    "# Rename columns (to deal with special characters)\n",
    "hrs_labels = ['IndID','FamID','BirthYear','HispanicStatus','DetailedHispanicStatus','Race_HRS','BirthRegionNum',\n",
    "             'BirthRegionName','AgeRange','Gender','Race_dbGaP']\n",
    "hrs_df = pd.DataFrame.from_records(aux_data[1:],columns = hrs_labels)\n",
    "hrs_df[['FamID','IndID']] = hrs_df[['FamID','IndID']].apply(pd.to_numeric)\n",
    "\n",
    "# Create a new dataset and reset the index\n",
    "hrs_df_sorted = hrs_df.sort_values(by=['FamID','IndID'])\n",
    "hrs_df_sorted = hrs_df_sorted.reset_index(drop=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "# Join the admixture data - this gets used at the end for different colourings.\n",
    "hrs_joined = hrs_df_sorted.merge(admix_df, left_on='IndID', right_on='ID', how='left')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "hrs_joined"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Get some counts\n",
    "#hrs_joined.groupby(['Race_HRS','HispanicStatus','DetailedHispanicStatus']).count()\n",
    "hrs_joined.groupby(['Race_HRS','HispanicStatus','DetailedHispanicStatus']).count()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "# Generalized version of the above\n",
    "# Create an auxiliary dataset for each subset\n",
    "# This is necessary so the indices match up when we want to attach a category/label to our points\n",
    "aux_data_dict = defaultdict(list)\n",
    "\n",
    "# Create a subset of our population and add it to the dict\n",
    "for hrs_pop in hrs_subsets_dict:\n",
    "    if hrs_pop in ['White','Black']:\n",
    "        hrs_df_subset = hrs_df_sorted.loc[(hrs_df_sorted['Race_HRS'] == hrs_pop)]\n",
    "    elif hrs_pop in ['Hispanic']:\n",
    "        hrs_df_subset = hrs_df_sorted.loc[(hrs_df_sorted['HispanicStatus'] == 'Hispanic')]\n",
    "    elif hrs_pop == 'All':\n",
    "        hrs_df_subset = hrs_df_sorted\n",
    "    \n",
    "    hrs_df_subset_list = hrs_df_subset.values.tolist()\n",
    "    aux_data_dict[hrs_pop] = hrs_df_subset_list\n",
    "    \n",
    "#aux_data_dict"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Subset the data here, if desired."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "try:\n",
    "    subset\n",
    "except NameError:\n",
    "    print('Subset not defined. Proceeding with full set of observations.')\n",
    "else:\n",
    "    if subset == 'White':\n",
    "        hrs_df_sorted = hrs_df_sorted.loc[(hrs_df_sorted['Race_HRS'] == 'White')]\n",
    "        print('Subsetting HRS data: White')\n",
    "    elif subset == 'Black':\n",
    "        hrs_df_sorted = hrs_df_sorted.loc[(hrs_df_sorted['Race_HRS'] == 'Black')]\n",
    "        print('Subsetting HRS data: Black')\n",
    "    elif subset == 'Hispanic':\n",
    "        hrs_df_sorted = hrs_df_sorted.loc[(hrs_df_sorted['HispanicStatus'] == 'Hispanic')]\n",
    "        print('Subsetting HRS data: Hispanic')\n",
    "    elif subset == 'All':\n",
    "        hrs_df_sorted = hrs_df_sorted\n",
    "        print('Using all entries - no subsetting carried out.')\n",
    "    else:\n",
    "        print('Subset \"' + str(subset) + '\" not recognized. Using full HRS dataset.')\n",
    "    \n",
    "# Use the sorted dataframe for auxiliary data\n",
    "aux_data_sorted = hrs_df_sorted.values.tolist()\n",
    "aux_data_sorted.insert(0,aux_data[0])\n",
    "\n",
    "aux_data = aux_data_sorted"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "# Define our dictionaries.\n",
    "# Create dicts for these values.\n",
    "hisp_dict = {'Hispanic':'H', 'Not_Hispanic':'N'}\n",
    "mex_dict = {'Mexican-American':'M', 'N/A':'N', 'Other':'O', 'Type_Unknown':'U'}\n",
    "race_dict = {'Black':'B', 'Other':'O', 'White':'W'}\n",
    "brn_dict = {'East_North_Central':'ENC', 'East_South_Central':'ESC', 'Middle_Atlantic':'MAT',\n",
    "           'Mountain':'MNT', 'New_England':'ENG', 'Not_In_Contiguous_US':'NIC','Pacific':'PAC',\n",
    "            'South_Atlantic':'SAT', 'West_North_Central':'WNC', 'West_South_Central':'WSC'}\n",
    "racedb_dict = {'AfrAm':'AA', 'Not_AfrAm':'NAA'}\n",
    "\n",
    "# Create reverse lookups. We don't have to use defaultdicts as this is 1-1\n",
    "hisp_dict_rev = dict()\n",
    "mex_dict_rev = dict()\n",
    "race_dict_rev = dict()\n",
    "brn_dict_rev = dict()\n",
    "racedb_dict_rev = dict()\n",
    "\n",
    "for key, value in hisp_dict.items():\n",
    "    hisp_dict_rev.update({value: key})\n",
    "    \n",
    "for key, value in mex_dict.items():\n",
    "    mex_dict_rev.update({value: key})\n",
    "\n",
    "for key, value in race_dict.items():\n",
    "    race_dict_rev.update({value: key})\n",
    "\n",
    "for key, value in brn_dict.items():\n",
    "    brn_dict_rev.update({value: key})\n",
    "    \n",
    "for key, value in racedb_dict.items():\n",
    "    racedb_dict_rev.update({value: key})"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Sharing variables between notebooks\n",
    "aux_data_hrs = aux_data\n",
    "%store aux_data_hrs\n",
    "%store hisp_dict\n",
    "%store hisp_dict_rev\n",
    "%store mex_dict\n",
    "%store mex_dict_rev\n",
    "%store race_dict\n",
    "%store race_dict_rev\n",
    "%store brn_dict\n",
    "%store brn_dict_rev\n",
    "%store racedb_dict\n",
    "%store racedb_dict_rev"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "# These auxiliary data sets will define how we label the observations\n",
    "\n",
    "# Columns are:\n",
    "# 0 = ID, 1 = Family ID, 2 = Birth Year\n",
    "# 3 = Hispanic, 4 = Detailed Hispanic, 5 = Race, 6 = Birth Region, 7 = Birth region name\n",
    "# 10= dbGaP race (Note: Black != AfrAm and White != Not_AfrAm)\n",
    "\n",
    "# Create multiple types of categorization based on variables to include\n",
    "aux_data_1 = [] # 1 - Birth region, race, Hispanic status, Mexican status\n",
    "aux_data_2 = [] # 2 - Race, Hispanic status, Mexican status\n",
    "aux_data_3 = [] # 3 - Birth region, race\n",
    "aux_data_4 = [] # 4 - Race, Hispanic status\n",
    "aux_data_5 = [] # 5 - Birth region\n",
    "aux_data_6 = [] # 6 - Birth region, Hispanic status, Mexican status\n",
    "\n",
    "individuals_hrs = []\n",
    "\n",
    "# Get the lists (skip the first row as it's a header)\n",
    "#for a in aux_data[1:]:\n",
    "for a in aux_data_dict[subset][0:]:\n",
    "    individuals_hrs.append(a[0])\n",
    "    \n",
    "    temp_element = [a[0], '_'.join([brn_dict[a[7]], race_dict[a[5]], hisp_dict[a[3]], mex_dict[a[4]]])]    \n",
    "    aux_data_1.append(temp_element)\n",
    "    \n",
    "    temp_element = [a[0], '_'.join([race_dict[a[5]], hisp_dict[a[3]], mex_dict[a[4]]])]\n",
    "    aux_data_2.append(temp_element)\n",
    "    \n",
    "    temp_element = [a[0], '_'.join([brn_dict[a[7]], race_dict[a[5]]])]\n",
    "    aux_data_3.append(temp_element)\n",
    "    \n",
    "    temp_element = [a[0], '_'.join([race_dict[a[5]], hisp_dict[a[3]]])]\n",
    "    aux_data_4.append(temp_element)\n",
    "    \n",
    "    temp_element = [a[0], '_'.join([brn_dict[a[7]]])]\n",
    "    aux_data_5.append(temp_element)\n",
    "    \n",
    "    temp_element = [a[0], '_'.join([brn_dict[a[7]], hisp_dict[a[3]], mex_dict[a[4]]])]\n",
    "    aux_data_6.append(temp_element)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "#%store aux_data_1\n",
    "#%store aux_data_2\n",
    "#%store aux_data_3\n",
    "#%store aux_data_4\n",
    "#%store aux_data_5\n",
    "#%store aux_data_6\n",
    "\n",
    "#%store individuals_hrs"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "# To save time, we decide which dataset we use here\n",
    "aux_to_use = aux_data_5"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "# We must define the population dictionary we wish to use\n",
    "# The following gives us a collection of all categories of some population and/or proxy for ethnicity:\n",
    "eth_proxy_set = set([a[1] for a in aux_to_use])\n",
    "pop_dict = dict()\n",
    "\n",
    "for e in eth_proxy_set:\n",
    "    el = e.split('_')\n",
    "    \n",
    "    if aux_to_use == aux_data_1:    \n",
    "        temp_brn = brn_dict_rev[el[0]]\n",
    "        temp_race = race_dict_rev[el[1]]\n",
    "        temp_hisp = hisp_dict_rev[el[2]]\n",
    "        temp_mex = mex_dict_rev[el[3]]\n",
    "        \n",
    "        pop_dict.update({e:temp_brn + ' ' + temp_race + ' ' + temp_hisp + ' ' + temp_mex})\n",
    "    elif aux_to_use == aux_data_2:\n",
    "        temp_race = race_dict_rev[el[0]]\n",
    "        temp_hisp = hisp_dict_rev[el[1]]\n",
    "        temp_mex = mex_dict_rev[el[2]]\n",
    "        temp_list = [temp_race, temp_hisp, temp_mex]\n",
    "        \n",
    "        pop_dict.update({e:temp_race + ' ' + temp_hisp + ' ' + temp_mex})\n",
    "    elif aux_to_use == aux_data_3:\n",
    "        temp_brn = brn_dict_rev[el[0]]\n",
    "        temp_race = race_dict_rev[el[1]]\n",
    "        \n",
    "        pop_dict.update({e: temp_brn + ' ' + temp_race})\n",
    "    elif aux_to_use == aux_data_4:\n",
    "        temp_race = race_dict_rev[el[0]]\n",
    "        temp_hisp = hisp_dict_rev[el[1]]\n",
    "        \n",
    "        pop_dict.update({e: temp_race + ' ' + temp_hisp})\n",
    "    elif aux_to_use == aux_data_5:\n",
    "        temp_brn = brn_dict_rev[el[0]]\n",
    "        \n",
    "        pop_dict.update({e: temp_brn})\n",
    "    elif aux_to_use == aux_data_6:\n",
    "        temp_brn = brn_dict_rev[el[0]]\n",
    "        temp_hisp = hisp_dict_rev[el[1]]\n",
    "        temp_mex = mex_dict_rev[el[2]]\n",
    "        \n",
    "        pop_dict.update({e: temp_brn + ' ' + temp_hisp + ' ' + temp_mex})\n",
    "    \n",
    "    #print(el, temp_list)\n",
    "    #print(e.split('_'))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "pop_dict"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "# Define a colour set - this supports up to 30 colours\n",
    "from bokeh.palettes import Category20\n",
    "from bokeh.palettes import PRGn\n",
    "from bokeh.palettes import Set1\n",
    "color_dict_hrs = {}\n",
    "\n",
    "for j, pop in enumerate(eth_proxy_set):\n",
    "    if j < 20:\n",
    "        color_dict_hrs[pop] = Category20[20][j]\n",
    "    elif j < 30:\n",
    "        color_dict_hrs[pop] = PRGn[10][j%20]\n",
    "    else:\n",
    "        color_dict_hrs[pop] = Set1[9][j%30]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "from bokeh.palettes import Category20c\n",
    "# The code from above works okay but sometimes generates different colours for same groups (i.e. inconsistent colours\n",
    "# between our plots)\n",
    "# Define some more specific colour dictionaries for the following three categories:\n",
    "# 1 - Birth region (10)\n",
    "# 2 - Race + Hispanic status (6)\n",
    "# 3 - Race + Hispanic status + Mexican status (8)\n",
    "\n",
    "# Birth regions are US census regions. Five divisions:\n",
    "# WEST: Pacfic (PAC), Mountain (MNT)\n",
    "# MIDWEST: West North Central (WNC), East North Central (ENC)\n",
    "# SOUTH: West South Central (WSC), East South Central (ESC), South Atlantic (SAT)\n",
    "# NORTHEAST: Middle Atlantic (MAT), New England (ENG)\n",
    "# Not in Contiguous US (NIC)\n",
    "\n",
    "color_dict_born = {}\n",
    "\n",
    "color_dict_born['ENG']=Category20b[20][1] # New England (purple)\n",
    "color_dict_born['MAT']=Category20b[20][3] # Mid Atlantic (light purple)\n",
    "color_dict_born['SAT']=Category20b[20][-2] # South Atlantic (pinkish)\n",
    "color_dict_born['ESC']=Category20b[20][-4] # East South Central (purplish-pink)\n",
    "color_dict_born['WSC']=Category20b[20][-6] # West South Central (rose-ish)\n",
    "color_dict_born['ENC']=Category20c[20][3] # East North Central (light blue)\n",
    "color_dict_born['WNC']=Category20c[20][0] # West North Central (blue)\n",
    "color_dict_born['MNT']=Category20b[20][5] # Mountain (Green)\n",
    "color_dict_born['PAC']=Category20b[20][7] # Pacific (lighter green)\n",
    "color_dict_born['NIC']=Category20c[20][-3] # Not in contiguous US (grey)\n",
    "\n",
    "color_dict_race_hisp = {}\n",
    "\n",
    "color_dict_race_hisp['B_H']=Category20b[20][2] # Black and Hispanic\n",
    "color_dict_race_hisp['B_N']=Category20b[20][3] # Black and not Hispanic\n",
    "color_dict_race_hisp['O_H']=Category20b[20][4] # Other and Hispanic\n",
    "color_dict_race_hisp['O_N']=Category20b[20][6] # Other and not Hispanic\n",
    "color_dict_race_hisp['W_H']=Category20c[4][-3] # White and Hispanic\n",
    "color_dict_race_hisp['W_N']=Category20c[4][-1] # White and not Hispanic\n",
    "\n",
    "color_dict_race_hisp_mex = {}\n",
    "\n",
    "color_dict_race_hisp_mex['B_H_O']=Category20b[20][2]# Black Hispanic non-Mexican (purple)\n",
    "color_dict_race_hisp_mex['B_N_N']=Category20b[20][3] # Black not Hispanic (light purple)\n",
    "color_dict_race_hisp_mex['O_H_M']=Category20b[20][4] # Other Hispanic Mexican (dark green)\n",
    "color_dict_race_hisp_mex['O_H_O']=Category20b[20][6] # Other Hispanic non-Mexican (green)\n",
    "color_dict_race_hisp_mex['O_N_N']=Category20b[20][-1] # Other not Hispanic non-Mexican (pink)\n",
    "color_dict_race_hisp_mex['W_H_M']=Category20c[4][-1] # White Hispanic Mexican (light blue)\n",
    "color_dict_race_hisp_mex['W_H_O']=Category20c[4][-2] # White Hispanic non-Mexican (less light blue)\n",
    "color_dict_race_hisp_mex['W_H_U']=Category20c[4][-3] # White Hispanic Unknown (even less light blue)\n",
    "color_dict_race_hisp_mex['W_N_N']=Category20c[4][0] # White not Hispanic non-Mexican (blue!)\n",
    "\n",
    "# Choose less terrible colours (but keep above for posterity)\n",
    "color_dict_race_hisp_mex['B_H_O']=Category20c[5][-1]# Black Hispanic non-Mexican (orange)\n",
    "color_dict_race_hisp_mex['B_N_N']=Category20c[7][-1] # Black not Hispanic (light orange)\n",
    "color_dict_race_hisp_mex['O_H_M']=Category20b[20][4] # Other Hispanic Mexican (dark green)\n",
    "color_dict_race_hisp_mex['O_H_O']=Category20b[20][6] # Other Hispanic non-Mexican (green)\n",
    "color_dict_race_hisp_mex['O_N_N']=Category20c[13][-1] # Other not Hispanic non-Mexican (poiple)\n",
    "color_dict_race_hisp_mex['W_H_M']=Category20c[9][-1] # White Hispanic Mexican (green)\n",
    "color_dict_race_hisp_mex['W_H_O']=Category20c[11][-1] # White Hispanic non-Mexican (lighter green)\n",
    "color_dict_race_hisp_mex['W_H_U']=Category20c[12][-1] # White Hispanic Unknown (even lighter green)\n",
    "color_dict_race_hisp_mex['W_N_N']=Category20c[4][0] # White not Hispanic non-Mexican (blue!)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%store color_dict_born\n",
    "%store color_dict_race_hisp\n",
    "%store color_dict_race_hisp_mex"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "# Set up an index of each population member and vice versa\n",
    "# We want to quickly access a given individual's population and a given population's individuals\n",
    "population_by_individual_hrs = defaultdict(int)\n",
    "individuals_by_population_hrs = defaultdict(list)\n",
    "\n",
    "for a in aux_to_use:\n",
    "    population_by_individual_hrs[a[0]] = a[1]\n",
    "    individuals_by_population_hrs[a[1]].append(a[0])\n",
    "    \n",
    "indices_of_population_members_hrs = defaultdict(list)\n",
    "\n",
    "for index, indiv in enumerate(individuals_hrs):\n",
    "    try:\n",
    "        indices_of_population_members_hrs[population_by_individual_hrs[indiv]].append(index)\n",
    "    except KeyError:\n",
    "        continue"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Colouring function\n",
    "Function to do all of the above more quickly (just specify the desired colouring)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "# Define a colour set - this supports up to 30 colours\n",
    "from bokeh.palettes import Category20\n",
    "from bokeh.palettes import PRGn\n",
    "from bokeh.palettes import Set1\n",
    "\n",
    "# These auxiliary data sets will define how we label the observations\n",
    "\n",
    "# Columns are:\n",
    "# 0 = ID, 1 = Family ID, 2 = Birth Year\n",
    "# 3 = Hispanic, 4 = Detailed Hispanic, 5 = Race, 6 = Birth Region, 7 = Birth region name\n",
    "# 10= dbGaP race (Note: Black != AfrAm and White != Not_AfrAm)\n",
    "\n",
    "# Create multiple types of categorization based on variables to include\n",
    "# 1 - Birth region, race, Hispanic status, Mexican status\n",
    "# 2 - Race, Hispanic status, Mexican status\n",
    "# 3 - Birth region, race\n",
    "# 4 - Race, Hispanic status\n",
    "# 5 - Birth region\n",
    "# 6 - Birth region, Hispanic status, Mexican status\n",
    "\n",
    "# We must define the population dictionary we wish to use\n",
    "# The following gives us a collection of all categories of some population and/or proxy for ethnicity:\n",
    "eth_proxy_set = set([a[1] for a in aux_to_use])\n",
    "pop_dict = dict()\n",
    "\n",
    "for e in eth_proxy_set:\n",
    "    el = e.split('_')\n",
    "    \n",
    "    if aux_to_use == aux_data_1:    \n",
    "        temp_brn = brn_dict_rev[el[0]]\n",
    "        temp_race = race_dict_rev[el[1]]\n",
    "        temp_hisp = hisp_dict_rev[el[2]]\n",
    "        temp_mex = mex_dict_rev[el[3]]\n",
    "        \n",
    "        pop_dict.update({e:temp_brn + ' ' + temp_race + ' ' + temp_hisp + ' ' + temp_mex})\n",
    "    elif aux_to_use == aux_data_2:\n",
    "        temp_race = race_dict_rev[el[0]]\n",
    "        temp_hisp = hisp_dict_rev[el[1]]\n",
    "        temp_mex = mex_dict_rev[el[2]]\n",
    "        temp_list = [temp_race, temp_hisp, temp_mex]\n",
    "        \n",
    "        pop_dict.update({e:temp_race + ' ' + temp_hisp + ' ' + temp_mex})\n",
    "    elif aux_to_use == aux_data_3:\n",
    "        temp_brn = brn_dict_rev[el[0]]\n",
    "        temp_race = race_dict_rev[el[1]]\n",
    "        \n",
    "        pop_dict.update({e: temp_brn + ' ' + temp_race})\n",
    "    elif aux_to_use == aux_data_4:\n",
    "        temp_race = race_dict_rev[el[0]]\n",
    "        temp_hisp = hisp_dict_rev[el[1]]\n",
    "        \n",
    "        pop_dict.update({e: temp_race + ' ' + temp_hisp})\n",
    "    elif aux_to_use == aux_data_5:\n",
    "        temp_brn = brn_dict_rev[el[0]]\n",
    "        \n",
    "        pop_dict.update({e: temp_brn})\n",
    "    elif aux_to_use == aux_data_6:\n",
    "        temp_brn = brn_dict_rev[el[0]]\n",
    "        temp_hisp = hisp_dict_rev[el[1]]\n",
    "        temp_mex = mex_dict_rev[el[2]]\n",
    "        \n",
    "        pop_dict.update({e: temp_brn + ' ' + temp_hisp + ' ' + temp_mex})\n",
    "\n",
    "color_dict_hrs = {}\n",
    "\n",
    "for j, pop in enumerate(eth_proxy_set):\n",
    "    if j < 20:\n",
    "        color_dict_hrs[pop] = Category20[20][j]\n",
    "    elif j < 30:\n",
    "        color_dict_hrs[pop] = PRGn[10][j%20]\n",
    "    else:\n",
    "        color_dict_hrs[pop] = Set1[9][j%30]\n",
    "\n",
    "# Set up an index of each population member and vice versa\n",
    "# We want to quickly access a given individual's population and a given population's individuals\n",
    "population_by_individual_hrs = defaultdict(int)\n",
    "individuals_by_population_hrs = defaultdict(list)\n",
    "\n",
    "for a in aux_to_use:\n",
    "    population_by_individual_hrs[a[0]] = a[1]\n",
    "    individuals_by_population_hrs[a[1]].append(a[0])\n",
    "    \n",
    "indices_of_population_members_hrs = defaultdict(list)\n",
    "\n",
    "for index, indiv in enumerate(individuals_hrs):\n",
    "    try:\n",
    "        indices_of_population_members_hrs[population_by_individual_hrs[indiv]].append(index)\n",
    "    except KeyError:\n",
    "        continue"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Projections and plotting\n",
    "\n",
    "Here we actually carry out the projections and plot the results."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "color_dict_born, color_dict_hrs"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "# Define a function to generate interactive HTML files\n",
    "def hrs_create_int_html(proj, fig_title, fname, page_title):\n",
    "    # Import TSNE PC projections from file and export HTML\n",
    "    temp_array = proj\n",
    "\n",
    "    component_1_id = 0\n",
    "    component_2_id = 1\n",
    "\n",
    "    p = figure(plot_width=1500, plot_height=800)\n",
    "    p.title.text = fig_title\n",
    "\n",
    "    for pop in sorted(eth_proxy_set):\n",
    "        proj_pop = temp_array[indices_of_population_members_hrs[pop]]\n",
    "        p.circle(proj_pop[:,component_1_id], proj_pop[:,component_2_id], legend=pop_dict[pop],\n",
    "                 color = color_dict_born[pop])\n",
    "\n",
    "    p.legend.location = \"top_left\"\n",
    "\n",
    "    p.legend.click_policy=\"hide\"\n",
    "\n",
    "    output_file(fname + '.html',title=page_title)\n",
    "\n",
    "    save(p)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "temp_array = np.loadtxt('/Volumes/Stockage/alex/hrs/projections/HRS_UMAP_PC40_NN15_MD0.01_201832924744')\n",
    "\n",
    "component_1_id = 0\n",
    "component_2_id = 1\n",
    "\n",
    "p = figure(plot_width=1500, plot_height=800)\n",
    "p.title.text = 'test'\n",
    "\n",
    "for pop in sorted(eth_proxy_set):\n",
    "    proj_pop = temp_array[indices_of_population_members_hrs[pop]]\n",
    "    p.circle(proj_pop[:,component_1_id], proj_pop[:,component_2_id], legend=pop_dict[pop],\n",
    "             color = color_dict_hrs[pop])\n",
    "\n",
    "p.legend.location = \"top_left\"\n",
    "\n",
    "p.legend.click_policy=\"hide\"\n",
    "\n",
    "output_file('test_HRS_UMAP_colour.html',title='test')\n",
    "\n",
    "save(p)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "proj_dir = '/Volumes/Stockage/alex/hrs/projections'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "# Set up labels\n",
    "if aux_to_use == aux_data_1:\n",
    "    aux_label = 'BORN_RACE_HISP_MEX'\n",
    "elif aux_to_use == aux_data_2:\n",
    "    aux_label = 'RACE_HISP_MEX'\n",
    "elif aux_to_use == aux_data_3:\n",
    "    aux_label = 'BORN_RACE'\n",
    "elif aux_to_use == aux_data_4:\n",
    "    aux_label = 'RACE_HISP'\n",
    "elif aux_to_use == aux_data_5:\n",
    "    aux_label = 'BORN'\n",
    "elif aux_to_use == aux_data_6:\n",
    "    aux_label = 'BORN_HISP_MEX'\n",
    "\n",
    "# Set up output directory\n",
    "out_dir = '/Volumes/Stockage/alex/hrs/html'\n",
    "\n",
    "# Loop through projections and create interactive HTML files\n",
    "for file in os.listdir(proj_dir):\n",
    "    if subset in file:\n",
    "        temp_proj = np.loadtxt(os.path.join(proj_dir,file))\n",
    "        \n",
    "        # Get details about the projection (number of PCs, etc) for organization\n",
    "        num_pcs = file.split('PC_')[1].split('_')[0]\n",
    "        plex_val = file.split('PLEX_')[1]\n",
    "        \n",
    "        # Naming format is HRS_TSNE_[subset]_PC_(num PCs)_PLEX_(plex value).html\n",
    "        out_file = os.path.join(out_dir,'HRS_TSNE_'+subset+'_'+aux_label+'_PC_'+num_pcs+'_PLEX_'+plex_val)\n",
    "        out_fig_title = 'HRS TSNE PC ' + num_pcs + ' plex ' + plex_val + ' pop subset: ' + subset\n",
    "        out_page_title = 'HRS TSNE on ' + num_pcs + ' PCs, perplexity ' + plex_val\n",
    "        \n",
    "        hrs_create_int_html(temp_proj, out_fig_title, out_file, out_page_title)\n",
    "        print('Finished creating HTML for file ' + file)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Plotting with admixture data\n",
    "\n",
    "We can repeat the above plots but use admixture data to colour the points. Colours are currently added by populations so this will be a slight tweak. We still want to keep the population groups for perusal but need colours defined elsewhere...."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Define a colour dictionary based on individual admixture proportions.\n",
    "# We have three values - this works nicely with an RGB tuple.\n",
    "hrs_joined_sorted = hrs_joined.values.tolist()\n",
    "#hrs_joined_sorted[2][-3:]\n",
    "#hrs_joined_sorted_array = np.array(hrs_joined_sorted[:,-3:])\n",
    "temp_list = [h[-3:] for h in hrs_joined_sorted]\n",
    "hrs_joined_sorted_array = np.array(temp_list)\n",
    "hrs_joined_sorted_array = (255*hrs_joined_sorted_array).astype(np.int64)\n",
    "hrs_joined_sorted_array.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "color_dict_admix = defaultdict(int)\n",
    "for i in range(0, len(hrs_joined_sorted)):\n",
    "    color_dict_admix[i] = '#%02x%02x%02x' % (hrs_joined_sorted_array[i][0],\n",
    "                                             hrs_joined_sorted_array[i][1],\n",
    "                                             hrs_joined_sorted_array[i][2])\n",
    "\n",
    "color_list_admix = list()\n",
    "for i in range(0, len(hrs_joined_sorted)):\n",
    "    color_list_admix.append('#%02x%02x%02x' % (hrs_joined_sorted_array[i][0],\n",
    "                                             hrs_joined_sorted_array[i][1],\n",
    "                                             hrs_joined_sorted_array[i][2]))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "len(color_dict_admix)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "# Test the colours on an arbitrary projection\n",
    "temp_proj = np.loadtxt('//Volumes/Stockage/alex/hrs/projections/HRS_UMAP_PC15_NN15_MD0.5_2018330153123')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "temp_proj.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "# Import TSNE PC projections from file and export HTML\n",
    "component_1_id = 0\n",
    "component_2_id = 1\n",
    "\n",
    "p = figure(plot_width=1500, plot_height=800)\n",
    "p.title.text = 'test UMAP'\n",
    "\n",
    "for pop in sorted(eth_proxy_set):\n",
    "    proj_pop = temp_proj[indices_of_population_members_hrs[pop]]\n",
    "    p.circle(proj_pop[:,component_1_id], proj_pop[:,component_2_id], legend=pop_dict[pop],\n",
    "             color = [color_list_admix[i] for i in indices_of_population_members_hrs[pop]])\n",
    "\n",
    "p.legend.location = \"top_left\"\n",
    "\n",
    "p.legend.click_policy=\"hide\"\n",
    "\n",
    "output_file('test2' + '.html',title='test UMAP')\n",
    "\n",
    "save(p)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "def hrs_tsne_admix(proj, ptitle, fname, ftitle):\n",
    "    # Import TSNE PC projections from file and export HTML\n",
    "    p = figure(plot_width=1500, plot_height=800)\n",
    "    p.title.text = ptitle\n",
    "\n",
    "    for pop in sorted(eth_proxy_set):\n",
    "        proj_pop = proj[indices_of_population_members_hrs[pop]]\n",
    "        p.circle(proj_pop[:,0], proj_pop[:,1], legend=pop_dict[pop],\n",
    "                 color = [color_list_admix[i] for i in indices_of_population_members_hrs[pop]])\n",
    "\n",
    "    p.legend.location = \"top_left\"\n",
    "\n",
    "    p.legend.click_policy=\"hide\"\n",
    "\n",
    "    output_file(fname + '.html',title=ftitle)\n",
    "\n",
    "    save(p)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "hrs_tsne_dir = '/Volumes/Stockage/alex/hrs/projections'\n",
    "\n",
    "for file in os.listdir(hrs_tsne_dir):\n",
    "    num_pcs = file.split('_PC_')[1].split('_')[0]\n",
    "    plex_val = file.split('_PC_')[1].split('_')[2]\n",
    "    temp_proj = np.loadtxt(os.path.join(hrs_tsne_dir,file))\n",
    "    \n",
    "    ptitle_str = 'tSNE on HRS data, first ' + str(num_pcs) + ' PCs, perplexity ' + str(plex_val) + ' (admix colours)'\n",
    "    fname_str = 'HRS_TSNE_PC_' + str(num_pcs) + '_PLEX_' + str(plex_val) + '_ADMIX'\n",
    "    ftitle_str = 'tSNE on HRS, ' + str(num_pcs) + ' PCs, plex ' + str(plex_val) + ' (admix colours)'\n",
    "    \n",
    "    hrs_tsne_admix(temp_proj, ptitle_str, fname_str, ftitle_str)\n",
    "    print('Generated interactive HTML for tSNE projections of HRS with admix colours on {PC,PLEX}: ' + str(num_pcs) +\n",
    "         ',' + str(plex_val))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# General HTML generator (requires running some previous steps)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "# Define a colour set - this supports up to 30 colours\n",
    "from bokeh.palettes import Category20\n",
    "from bokeh.palettes import PRGn\n",
    "from bokeh.palettes import Set1"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "len(aux_data_5)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "aux_to_use = aux_data_5\n",
    "# These auxiliary data sets will define how we label the observations\n",
    "\n",
    "# Columns are:\n",
    "# 0 = ID, 1 = Family ID, 2 = Birth Year\n",
    "# 3 = Hispanic, 4 = Detailed Hispanic, 5 = Race, 6 = Birth Region, 7 = Birth region name\n",
    "# 10= dbGaP race (Note: Black != AfrAm and White != Not_AfrAm)\n",
    "\n",
    "# Create multiple types of categorization based on variables to include\n",
    "# 1 - Birth region, race, Hispanic status, Mexican status\n",
    "# 2 - Race, Hispanic status, Mexican status\n",
    "# 3 - Birth region, race\n",
    "# 4 - Race, Hispanic status\n",
    "# 5 - Birth region\n",
    "# 6 - Birth region, Hispanic status, Mexican status\n",
    "\n",
    "# We must define the population dictionary we wish to use\n",
    "# The following gives us a collection of all categories of some population and/or proxy for ethnicity:\n",
    "eth_proxy_set = set([a[1] for a in aux_to_use])\n",
    "pop_dict = dict()\n",
    "\n",
    "for e in eth_proxy_set:\n",
    "    el = e.split('_')\n",
    "    \n",
    "    if aux_to_use == aux_data_1:    \n",
    "        temp_brn = brn_dict_rev[el[0]]\n",
    "        temp_race = race_dict_rev[el[1]]\n",
    "        temp_hisp = hisp_dict_rev[el[2]]\n",
    "        temp_mex = mex_dict_rev[el[3]]\n",
    "        \n",
    "        pop_dict.update({e:temp_brn + ' ' + temp_race + ' ' + temp_hisp + ' ' + temp_mex})\n",
    "    elif aux_to_use == aux_data_2:\n",
    "        temp_race = race_dict_rev[el[0]]\n",
    "        temp_hisp = hisp_dict_rev[el[1]]\n",
    "        temp_mex = mex_dict_rev[el[2]]\n",
    "        temp_list = [temp_race, temp_hisp, temp_mex]\n",
    "        \n",
    "        pop_dict.update({e:temp_race + ' ' + temp_hisp + ' ' + temp_mex})\n",
    "    elif aux_to_use == aux_data_3:\n",
    "        temp_brn = brn_dict_rev[el[0]]\n",
    "        temp_race = race_dict_rev[el[1]]\n",
    "        \n",
    "        pop_dict.update({e: temp_brn + ' ' + temp_race})\n",
    "    elif aux_to_use == aux_data_4:\n",
    "        temp_race = race_dict_rev[el[0]]\n",
    "        temp_hisp = hisp_dict_rev[el[1]]\n",
    "        \n",
    "        pop_dict.update({e: temp_race + ' ' + temp_hisp})\n",
    "    elif aux_to_use == aux_data_5:\n",
    "        temp_brn = brn_dict_rev[el[0]]\n",
    "        \n",
    "        pop_dict.update({e: temp_brn})\n",
    "    elif aux_to_use == aux_data_6:\n",
    "        temp_brn = brn_dict_rev[el[0]]\n",
    "        temp_hisp = hisp_dict_rev[el[1]]\n",
    "        temp_mex = mex_dict_rev[el[2]]\n",
    "        \n",
    "        pop_dict.update({e: temp_brn + ' ' + temp_hisp + ' ' + temp_mex})\n",
    "\n",
    "color_dict_hrs = {}\n",
    "\n",
    "for j, pop in enumerate(eth_proxy_set):\n",
    "    if j < 20:\n",
    "        color_dict_hrs[pop] = Category20[20][j]\n",
    "    elif j < 30:\n",
    "        color_dict_hrs[pop] = PRGn[10][j%20]\n",
    "    else:\n",
    "        color_dict_hrs[pop] = Set1[9][j%30]\n",
    "\n",
    "# Set up an index of each population member and vice versa\n",
    "# We want to quickly access a given individual's population and a given population's individuals\n",
    "population_by_individual_hrs = defaultdict(int)\n",
    "individuals_by_population_hrs = defaultdict(list)\n",
    "\n",
    "for a in aux_to_use:\n",
    "    population_by_individual_hrs[a[0]] = a[1]\n",
    "    individuals_by_population_hrs[a[1]].append(a[0])\n",
    "    \n",
    "indices_of_population_members_hrs = defaultdict(list)\n",
    "\n",
    "for index, indiv in enumerate(individuals_hrs):\n",
    "    try:\n",
    "        indices_of_population_members_hrs[population_by_individual_hrs[indiv]].append(index)\n",
    "    except KeyError:\n",
    "        continue"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "proj_dir = '/Volumes/Stockage/alex/hrs/projections'\n",
    "out_dir = '/Volumes/Stockage/alex/hrs/html'\n",
    "\n",
    "subset = 'Hispanic'\n",
    "\n",
    "if aux_to_use == aux_data_1:\n",
    "    aux_label = 'BORN_RACE_HISP_MEX'\n",
    "elif aux_to_use == aux_data_2:\n",
    "    aux_label = 'RACE_HISP_MEX'\n",
    "elif aux_to_use == aux_data_3:\n",
    "    aux_label = 'BORN_RACE'\n",
    "elif aux_to_use == aux_data_4:\n",
    "    aux_label = 'RACE_HISP'\n",
    "elif aux_to_use == aux_data_5:\n",
    "    aux_label = 'BORN'\n",
    "elif aux_to_use == aux_data_6:\n",
    "    aux_label = 'BORN_HISP_MEX'\n",
    "\n",
    "for file in os.listdir(proj_dir):\n",
    "    \n",
    "    # Drop directories\n",
    "    if os.path.isdir(file):\n",
    "        continue\n",
    "        \n",
    "    # Skip already-created pages\n",
    "    if os.path.exists(os.path.join(out_dir, file + '_' + aux_label + '.html')) and \\\n",
    "    os.path.exists(os.path.join(out_dir, file + '_' + aux_label + '_ADMIX.html')):\n",
    "        continue\n",
    "    \n",
    "    if subset=='All':\n",
    "        # Only look for files related to this subset\n",
    "        if not ('HRS_UMAP' in file):\n",
    "            continue\n",
    "        try:\n",
    "            temp_proj = np.loadtxt(os.path.join(proj_dir, file))\n",
    "            out_fig_title = file\n",
    "            out_file = os.path.join(out_dir, file + '_' + aux_label)\n",
    "            out_page_title = file\n",
    "            hrs_create_int_html(temp_proj, out_fig_title, out_file, out_page_title)\n",
    "            hrs_tsne_admix(temp_proj, out_fig_title, out_file+'_ADMIX', out_page_title)\n",
    "        except Exception as e:\n",
    "            print(e)\n",
    "            print(file)\n",
    "    elif subset=='Hispanic':\n",
    "        # Work only with the Hispanic subset\n",
    "        if not ('HRS_HISP' in file):\n",
    "            continue\n",
    "        \n",
    "        try:\n",
    "            temp_proj = np.loadtxt(os.path.join(proj_dir, file))\n",
    "            out_fig_title = file\n",
    "            out_file = os.path.join(out_dir, file + '_' + aux_label)\n",
    "            out_page_title = file\n",
    "            hrs_create_int_html(temp_proj, out_fig_title, out_file, out_page_title)\n",
    "            #hrs_tsne_admix(temp_proj, out_fig_title, out_file+'_ADMIX', out_page_title)\n",
    "        except Exception as e:\n",
    "            print(e)\n",
    "            print(file)\n",
    "    elif subset=='Black':\n",
    "        # Work only with Black subset\n",
    "        if not ('HRS_BLACK' in file):\n",
    "            continue\n",
    "        \n",
    "        try:\n",
    "            temp_proj = np.loadtxt(os.path.join(proj_dir, file))\n",
    "            out_fig_title = file\n",
    "            out_file = os.path.join(out_dir, file + '_' + aux_label)\n",
    "            out_page_title = file\n",
    "            hrs_create_int_html(temp_proj, out_fig_title, out_file, out_page_title)\n",
    "            hrs_tsne_admix(temp_proj, out_fig_title, out_file+'_ADMIX', out_page_title)\n",
    "        except Exception as e:\n",
    "            print(e)\n",
    "            print(file)\n",
    "    elif subset=='White':\n",
    "        # Work only with White subset\n",
    "        if not ('HRS_WHITE' in file):\n",
    "            continue\n",
    "        \n",
    "        try:\n",
    "            temp_proj = np.loadtxt(os.path.join(proj_dir, file))\n",
    "            out_fig_title = file\n",
    "            out_file = os.path.join(out_dir, file + '_' + aux_label)\n",
    "            out_page_title = file\n",
    "            hrs_create_int_html(temp_proj, out_fig_title, out_file, out_page_title)\n",
    "            hrs_tsne_admix(temp_proj, out_fig_title, out_file+'_ADMIX', out_page_title)\n",
    "        except Exception as e:\n",
    "            print(e)\n",
    "            print(file)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Generate non-interactive images"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "color_dict_born"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true,
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "# Race-Hispanic-Mexican\n",
    "# Generate images (not HTML)\n",
    "proj_dir = '/Volumes/Stockage/alex/hrs/projections'\n",
    "out_dir = '/Volumes/Stockage/alex/hrs/images'\n",
    "\n",
    "x=0\n",
    "\n",
    "for fname in os.listdir(proj_dir):\n",
    "    if x > 1:\n",
    "        continue\n",
    "    elif os.path.isdir(os.path.join(proj_dir, fname)) or fname[0:8] not in ['HRS_TSNE','HRS_UMAP']:\n",
    "        continue\n",
    "    else:\n",
    "        proj = np.loadtxt(os.path.join(proj_dir, fname))\n",
    "\n",
    "        fig = plt.figure(figsize=(20,20))\n",
    "        ax = fig.add_subplot(111, aspect=1)\n",
    "\n",
    "        for pop in sorted(eth_proxy_set):\n",
    "            temp_proj = proj[indices_of_population_members_hrs[pop]]\n",
    "            ax.scatter(temp_proj[:,0], temp_proj[:,1], label=pop_dict[pop], alpha=0.6, color=color_dict_race_hisp_mex[pop])\n",
    "\n",
    "        ax.legend(ncol=3,loc='lower center', bbox_to_anchor=(0.55,-0.15), fontsize=12,markerscale=3)\n",
    "        fig.savefig(os.path.join(out_dir, fname + '.jpeg'),format='jpeg')\n",
    "        plt.close()\n",
    "\n",
    "        #x+=1"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "# Birth region\n",
    "# Generate images (not HTML)\n",
    "proj_dir = '/Volumes/Stockage/alex/hrs/projections'\n",
    "out_dir = '/Volumes/Stockage/alex/hrs/images'\n",
    "\n",
    "x=0\n",
    "\n",
    "for fname in os.listdir(proj_dir):\n",
    "    if x > 1:\n",
    "        continue\n",
    "    elif os.path.isdir(os.path.join(proj_dir, fname)) or fname[0:8] not in ['HRS_TSNE','HRS_UMAP']:\n",
    "        continue\n",
    "    else:\n",
    "        proj = np.loadtxt(os.path.join(proj_dir, fname))\n",
    "\n",
    "        fig = plt.figure(figsize=(20,20))\n",
    "        ax = fig.add_subplot(111, aspect=1)\n",
    "\n",
    "        for pop in sorted(eth_proxy_set):\n",
    "            temp_proj = proj[indices_of_population_members_hrs[pop]]\n",
    "            ax.scatter(temp_proj[:,0], temp_proj[:,1], label=pop_dict[pop], s=6, alpha=0.6, color=color_dict_born[pop])\n",
    "\n",
    "        ax.legend(ncol=3,loc='lower center', bbox_to_anchor=(0.55,-0.15), fontsize=12,markerscale=3)\n",
    "        fig.savefig(os.path.join(out_dir, fname + '_born.jpeg'),format='jpeg')\n",
    "        plt.close()\n",
    "\n",
    "        #x+=1"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Admixed populations"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "# Generate images (not HTML)\n",
    "proj_dir = '/Volumes/Stockage/alex/hrs/projections'\n",
    "out_dir = '/Volumes/Stockage/alex/hrs/images'\n",
    "\n",
    "x=0\n",
    "\n",
    "for fname in os.listdir(proj_dir):\n",
    "    if x > 1:\n",
    "        continue\n",
    "    elif os.path.isdir(os.path.join(proj_dir, fname)) or fname[0:8] not in ['HRS_TSNE','HRS_UMAP']:\n",
    "        continue\n",
    "    else:\n",
    "        proj = np.loadtxt(os.path.join(proj_dir, fname))\n",
    "\n",
    "        fig = plt.figure(figsize=(40,40))\n",
    "        ax = fig.add_subplot(111, aspect=1)\n",
    "\n",
    "        ax.scatter(proj[:,0], proj[:,1], alpha=0.6, c=color_list_admix)\n",
    "        \n",
    "        fig.savefig(os.path.join(out_dir, fname + '_admix.jpeg'),format='jpeg')\n",
    "        plt.close()\n",
    "\n",
    "        #x+=1"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "eth_proxy_set"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "# Hispanic subset, coloured by birth region\n",
    "proj_dir = '/Volumes/Stockage/alex/hrs/projections'\n",
    "out_dir = '/Volumes/Stockage/alex/hrs/images'\n",
    "\n",
    "fname = 'HRS_HISP_UMAP_PC10_NC2_NN15_MD0.5_201844193148'\n",
    "proj = np.loadtxt(os.path.join(proj_dir, fname))\n",
    "\n",
    "fig = plt.figure(figsize=(20,20))\n",
    "ax = fig.add_subplot(111, aspect=1)\n",
    "\n",
    "for pop in sorted(eth_proxy_set):\n",
    "    temp_proj = proj[indices_of_population_members_hrs[pop]]\n",
    "    ax.scatter(temp_proj[:,0], temp_proj[:,1], label=pop_dict[pop], alpha=0.6, color=color_dict_born[pop])\n",
    "\n",
    "ax.legend(ncol=3,loc='lower center', bbox_to_anchor=(0.55,-0.15), fontsize=12,markerscale=3)\n",
    "fig.savefig(os.path.join(out_dir, fname + '.jpeg'),format='jpeg')\n",
    "plt.close()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "FIX HERE (just need to line up indices for subsets)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "color_list_admix"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "len(color_list_admix)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "for k in indices_of_population_members_hrs.keys():\n",
    "    print(len(indices_of_population_members_hrs[k]))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "# Get indices of Hispanic individuals only and pare down our admixture colour list\n",
    "hispanic_index_list = hrs_joined[hrs_joined['HispanicStatus']=='Hispanic'].index.tolist()\n",
    "color_list_admix_hispanic = [color_list_admix[h] for h in hispanic_index_list]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true,
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "# Hispanic subset, coloured by admixture (probably??)\n",
    "proj_dir = '/Volumes/Stockage/alex/hrs/projections'\n",
    "out_dir = '/Volumes/Stockage/alex/hrs/images'\n",
    "\n",
    "fname = 'HRS_HISP_UMAP_PC7_NC2_NN15_MD0.5_201844193148'\n",
    "proj = np.loadtxt(os.path.join(proj_dir, fname))\n",
    "\n",
    "fig = plt.figure(figsize=(20,20))\n",
    "ax = fig.add_subplot(111, aspect=1)\n",
    "\n",
    "for pop in sorted(eth_proxy_set):\n",
    "    temp_proj = proj[indices_of_population_members_hrs[pop]]\n",
    "    color_list = [color_list_admix_hispanic[i] for i in indices_of_population_members_hrs[pop]]\n",
    "    ax.scatter(temp_proj[:,0], temp_proj[:,1], label=pop_dict[pop], alpha=0.6, color=color_list)\n",
    "\n",
    "ax.legend(ncol=3,loc='lower center', bbox_to_anchor=(0.55,-0.15), fontsize=12,markerscale=3)\n",
    "fig.savefig(os.path.join(out_dir, fname + '_admix.jpeg'),format='jpeg')\n",
    "plt.close()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "pop_dict"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "proj_dir = '/Volumes/Stockage/alex/hrs/projections'\n",
    "out_dir = '/Volumes/Stockage/alex/hrs/images'\n",
    "\n",
    "fname = 'HRS_HISP_PC1_PC2'\n",
    "proj = hrs_data_dict['Hispanic']\n",
    "fig = plt.figure(figsize=(20,20))\n",
    "ax = fig.add_subplot(111, aspect=1)\n",
    "\n",
    "for pop in sorted(eth_proxy_set):\n",
    "    temp_proj = proj[indices_of_population_members_hrs[pop]]\n",
    "    ax.scatter(temp_proj[:,0], temp_proj[:,1], label=pop_dict[pop], alpha=0.6, color=color_dict_race_hisp_mex[pop])\n",
    "\n",
    "ax.legend(ncol=3,loc='lower center', bbox_to_anchor=(0.55,-0.15), fontsize=12,markerscale=3)\n",
    "fig.savefig(os.path.join(out_dir, fname + '_RACE_HISP_MEX.jpeg'),format='jpeg')\n",
    "plt.close()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "proj_dir = '/Volumes/Stockage/alex/hrs/projections'\n",
    "out_dir = '/Volumes/Stockage/alex/hrs/images'\n",
    "\n",
    "fname = 'HRS_HISP_PC1_PC2'\n",
    "proj = hrs_data_dict['Hispanic']\n",
    "fig = plt.figure(figsize=(20,20))\n",
    "ax = fig.add_subplot(111, aspect=1)\n",
    "\n",
    "for pop in sorted(eth_proxy_set):\n",
    "    temp_proj = proj[indices_of_population_members_hrs[pop]]\n",
    "    ax.scatter(temp_proj[:,0], temp_proj[:,1], label=pop_dict[pop], alpha=0.6, color=color_dict_born[pop])\n",
    "\n",
    "ax.legend(ncol=3,loc='lower center', bbox_to_anchor=(0.55,-0.15), fontsize=12,markerscale=3)\n",
    "fig.savefig(os.path.join(out_dir, fname + '_BORN.jpeg'),format='jpeg')\n",
    "plt.close()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "#HRS_UMAP_PC10_NN15_MD0.5_2018330153123\n",
    "proj_dir = '/Volumes/Stockage/alex/hrs/projections'\n",
    "out_dir = '/Volumes/Stockage/alex/hrs/images/other'\n",
    "\n",
    "fname = 'HRS_UMAP_PC10_NN15_MD0.5_2018330153123'\n",
    "proj= np.loadtxt(os.path.join(proj_dir, fname))\n",
    "fig = plt.figure(figsize=(20,20))\n",
    "ax = fig.add_subplot(111, aspect=1)\n",
    "\n",
    "#for pop in sorted(eth_proxy_set):\n",
    "for pop in ['W_N_N','B_N_N','W_H_M','W_H_O','O_N_N','O_H_M','O_H_O','B_H_O','W_H_U']:\n",
    "    temp_proj = proj[indices_of_population_members_hrs[pop]]\n",
    "    ax.scatter(temp_proj[:,0], temp_proj[:,1], label=pop_dict[pop], alpha=0.6, color=color_dict_race_hisp_mex[pop])\n",
    "\n",
    "ax.legend(ncol=3,loc='lower center', bbox_to_anchor=(0.55,-0.15), fontsize=12,markerscale=3)\n",
    "fig.savefig(os.path.join(out_dir, fname + '_RACE_HISP_MEX.jpeg'),format='jpeg')\n",
    "plt.close()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "for e in eth_proxy_set:\n",
    "    print(e, proj[indices_of_population_members_hrs[e]].shape[0])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Admixture plots for Hispanic population\n",
    "Goal: Create an admixture bar plot.\n",
    "\n",
    "Need: Admixture for Hispanics (somewhere in here....) and indices and birth regions. Cluster or highlight mountain region Hispanic individuals"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "# Admixture data is stored in hrs_joined\n",
    "\n",
    "hrs_joined_hispanic = hrs_joined.loc[hrs_joined['HispanicStatus']=='Hispanic']\n",
    "hrs_joined_hispanic=hrs_joined_hispanic.sort_values(['BirthRegionName','ADMIX3'])\n",
    "hrs_joined_hispanic.reset_index(drop=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "afr_bars = [i for i in hrs_joined_hispanic['ADMIX1']]\n",
    "eur_bars = [i for i in hrs_joined_hispanic['ADMIX2']]\n",
    "oth_bars = [i for i in hrs_joined_hispanic['ADMIX3']]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "np.where(hrs_joined_hispanic['BirthRegionName']=='Mountain')[0]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "bar_width = 1.0\n",
    "inds = [i for i in range(0, hrs_joined_hispanic.count()[0])]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "plt.figure(figsize=(20,10))\n",
    "plt.bar(inds, oth_bars, color='blue', width=bar_width)\n",
    "plt.bar(inds, eur_bars, bottom = oth_bars, color='green', width=bar_width)\n",
    "#plt.bar(inds, eur_bars, bottom = [i+j for i,j in zip(afr_bars, oth_bars)], color='green', width=bar_width)\n",
    "plt.bar(inds, afr_bars, bottom = [i+j for i,j in zip(oth_bars, eur_bars)],color='red', width=bar_width)\n",
    "plt.xticks([0,48,184,400,600,800,1000,1200])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "# We can take this one step further. Use HRS_HISP_UMAP_PC10_NC2_NN15_MD0.5_201844193148 and find all instances of y<-10\n",
    "proj_dir = '/Volumes/Stockage/alex/hrs/projections'\n",
    "fname = 'HRS_HISP_UMAP_PC10_NC2_NN15_MD0.5_201844193148'\n",
    "\n",
    "proj= np.loadtxt(os.path.join(proj_dir, fname))\n",
    "proj_df = pd.DataFrame(proj)\n",
    "proj_df.columns=['UMAP_x','UMAP_y']"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true,
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "# Reload the dataframe (sorting screwed it up)\n",
    "hrs_joined_hispanic = hrs_joined.loc[hrs_joined['HispanicStatus']=='Hispanic']\n",
    "hrs_joined_hispanic_umap = pd.concat([hrs_joined_hispanic.reset_index(drop=True), proj_df], axis=1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true,
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "hrs_joined_hispanic_umap=hrs_joined_hispanic_umap.sort_values(['BirthRegionName','ADMIX3'])\n",
    "hrs_joined_hispanic_umap=hrs_joined_hispanic_umap.reset_index(drop=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true,
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "cluster = hrs_joined_hispanic_umap.loc[(hrs_joined_hispanic_umap['BirthRegionName']=='Mountain') & (hrs_joined_hispanic_umap['UMAP_y'] < -10)]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "cluster"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "cluster.count()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "collapsed": true
   },
   "source": [
    "# Sandbox"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Plot admixture colours on ethnic subsets"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "hrs_joined_sorted = hrs_joined.values.tolist()\n",
    "temp_list = [h[-3:] for h in hrs_joined_sorted]\n",
    "hrs_joined_sorted_array = np.array(temp_list)\n",
    "hrs_joined_sorted_array = (255*hrs_joined_sorted_array).astype(np.int64)\n",
    "hrs_joined_sorted_array.shape\n",
    "\n",
    "color_dict_admix = defaultdict(int)\n",
    "for i in range(0, len(hrs_joined_sorted)):\n",
    "    color_dict_admix[i] = '#%02x%02x%02x' % (hrs_joined_sorted_array[i][0],\n",
    "                                             hrs_joined_sorted_array[i][1],\n",
    "                                             hrs_joined_sorted_array[i][2])\n",
    "\n",
    "color_list_admix = list()\n",
    "for i in range(0, len(hrs_joined_sorted)):\n",
    "    color_list_admix.append('#%02x%02x%02x' % (hrs_joined_sorted_array[i][0],\n",
    "                                             hrs_joined_sorted_array[i][1],\n",
    "                                             hrs_joined_sorted_array[i][2]))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "hrs_data_dict['Hispanic'].shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "hispanic_index_list = hrs_df_sorted.loc[(hrs_df_sorted['HispanicStatus'] == 'Hispanic')].index.values.tolist()\n",
    "color_list_hispanic = [color_list_admix[i] for i in hispanic_index_list]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "len(color_list_hispanic)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "# Import TSNE PC projections from file and export HTML\n",
    "temp_proj = np.loadtxt('/Volumes/Stockage/alex/hrs/projections/HRS_HISP_UMAP_PC10_NC2_NN15_MD0.5_201844193148')\n",
    "p = figure(plot_width=1500, plot_height=800)\n",
    "p.title.text = 'Test Hispanic admixture values'\n",
    "\n",
    "p.circle(temp_proj[:,0], temp_proj[:,1], color = color_list_hispanic, size=5)\n",
    "\n",
    "p.legend.location = \"top_left\"\n",
    "\n",
    "p.legend.click_policy=\"hide\"\n",
    "\n",
    "output_file('test2' + '.html',title='test')\n",
    "\n",
    "show(p)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Import TSNE PC projections from file and export HTML\n",
    "#temp_proj = np.loadtxt('/Volumes/Stockage/alex/hrs/projections/HRS_PCA_hispanic.eigenvec')\n",
    "temp_proj = hrs_data_dict['Hispanic']\n",
    "p = figure(plot_width=1500, plot_height=800)\n",
    "p.title.text = 'Test Hispanic PC values'\n",
    "\n",
    "p.circle(temp_proj[:,0], temp_proj[:,1], color = color_list_hispanic, size=5)\n",
    "\n",
    "p.legend.location = \"top_left\"\n",
    "\n",
    "p.legend.click_policy=\"hide\"\n",
    "\n",
    "output_file('test3' + '.html',title='test')\n",
    "\n",
    "show(p)\n",
    "\n",
    "temp_proj.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "len(color_list_admix)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "indices_of_population_members_hrs.keys()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#hrs_create_int_html(hrs_data_dict['Hispanic'], 'HRS Hispanic', of, 'HRS Hispanic PCs')\n",
    "\n",
    "temp_array = hrs_data_dict['Hispanic']\n",
    "\n",
    "component_1_id = 4\n",
    "component_2_id = 5\n",
    "\n",
    "p = figure(plot_width=1500, plot_height=800)\n",
    "p.title.text = 'HRS Hispanic PCs ' + str(component_1_id+1) + ' and ' + str(component_2_id+1)\n",
    "\n",
    "for pop in sorted(eth_proxy_set):\n",
    "    proj_pop = temp_array[indices_of_population_members_hrs[pop]]\n",
    "    p.circle(proj_pop[:,component_1_id], proj_pop[:,component_2_id], legend=pop_dict[pop],\n",
    "             color = color_dict_born[pop])\n",
    "\n",
    "p.legend.location = \"top_left\"\n",
    "\n",
    "p.legend.click_policy=\"hide\"\n",
    "\n",
    "fname = 'HRS_HISP_PC' + str(component_1_id+1) + '_' + str(component_2_id+1) + '_BORN'\n",
    "\n",
    "output_file(fname + '.html',title='HRS Hispanic PCs ' + str(component_1_id+1) + ' and ' + str(component_2_id+1))\n",
    "\n",
    "save(p)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "temp_proj = np.loadtxt('/Volumes/Stockage/alex/hrs/projections/HRS_HISP_UMAP_PC10_NC2_NN15_MD0.5_201844193148')\n",
    "of = 'Volumes/Stockage/alex/hrs/html/HRS_HISP_UMAP_PC10_NC2_NN15_MD0.5_201844193148'\n",
    "hrs_create_int_html(temp_proj, 'HRS Hispanic', of, 'HRS Hispanic PCA-UMAP')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Create a grid of PC groups to look at Hispanic clustering"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "proj = hrs_data_dict['Hispanic']\n",
    "\n",
    "fig = plt.figure(figsize=(20,20))\n",
    "ax = fig.add_subplot(111, aspect=1)\n",
    "\n",
    "for pop in sorted(eth_proxy_set):\n",
    "    temp_proj = proj[indices_of_population_members_hrs[pop]]\n",
    "    ax.scatter(temp_proj[:,0], temp_proj[:,1], label=pop_dict[pop], alpha=0.6, color=color_dict_born[pop])\n",
    "\n",
    "ax.legend(ncol=3,loc='lower center', bbox_to_anchor=(0.55,-0.15), fontsize=12,markerscale=3)\n",
    "fig.savefig(os.path.join(out_dir, fname + '.jpeg'),format='jpeg')\n",
    "plt.close()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "proj = hrs_data_dict['Hispanic']\n",
    "\n",
    "gridsize = 7\n",
    "f, axarr = plt.subplots(gridsize,gridsize,figsize=(40,40))\n",
    "\n",
    "for i in range(0,gridsize):\n",
    "    for j in range(0,gridsize):\n",
    "        if i<=j:\n",
    "            for pop in eth_proxy_set:\n",
    "                temp_proj = proj[indices_of_population_members_hrs[pop]]\n",
    "                axarr[i,j].scatter(temp_proj[:,i], temp_proj[:,j+1], alpha=0.6, s=14, color=color_dict_born[pop])\n",
    "                axarr[i,j].set_title('PC' + str(j+2) + ' vs PC' + str(i+1))\n",
    "            \n",
    "    \n",
    "plt.legend(ncol=3,loc='lower center', bbox_to_anchor=[-1,-1], bbox_transform=plt.gcf().transFigure, fontsize=12,markerscale=3)   \n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "# Generate images (not HTML)\n",
    "proj_dir = '/Volumes/Stockage/alex/hrs/projections'\n",
    "out_dir = '/Volumes/Stockage/alex/hrs/images/grid'\n",
    "\n",
    "fname = 'HRS_HISP_UMAP_PC7_NC2_NN15_MD0.5_201844193148'\n",
    "\n",
    "proj = np.loadtxt(os.path.join(proj_dir, fname))\n",
    "\n",
    "fig = plt.figure(figsize=(20,20))\n",
    "ax = fig.add_subplot(111, aspect=1)\n",
    "\n",
    "for pop in sorted(eth_proxy_set):\n",
    "    temp_proj = proj[indices_of_population_members_hrs[pop]]\n",
    "    ax.scatter(temp_proj[:,0], temp_proj[:,1], label=pop_dict[pop], alpha=0.8, color=color_dict_born[pop])\n",
    "\n",
    "ax.legend(ncol=3,loc='lower center', bbox_to_anchor=(0.55,-0.1), fontsize=12,markerscale=3)\n",
    "fig.savefig(os.path.join(out_dir, fname + '.jpeg'),format='jpeg')\n",
    "plt.close()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "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.6.3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}