Data-Processing.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "0",
   "metadata": {},
   "source": [
    "# Data Preparation\n",
    "\n",
    "This notebook contains the code to generate the base datasets for later estimation. Steps include;\n",
    "\n",
    "   1. Removing people who do not meet the inclusion criteria (age at wave1, present in both intial and final wave). \n",
    "   2. Computing treatment and outcome variables.\n",
    "   3. Removing features that are deemed to be proxies for the treatment variable\n",
    "   4. Removing columns that are ids or otherwise deemed irrelevant.\n",
    "   5. Optional unsupervised feature selection (with respect to the treatment/target) feature selection to produce smaller datasets and reduce the issue of correlated features."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1",
   "metadata": {
    "tags": [
     "hide_cell"
    ]
   },
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import pandas as pd\n",
    "import matplotlib.pyplot as plt\n",
    "import pyreadstat\n",
    "import re\n",
    "import string\n",
    "#from sklearn_pandas import DataFrameMapper\n",
    "import networkx as nx\n",
    "import reed, config\n",
    "import pickle\n",
    "from clean import *\n",
    "from reed import regex_select\n",
    "from hilda_config import release_suffix\n",
    "\n",
    "pd.options.display.max_columns=100\n",
    "pd.options.display.max_colwidth=200\n",
    "%load_ext autoreload\n",
    "%autoreload 2\n",
    "\n",
    "from platform import python_version\n",
    "\n",
    "print(python_version())\n",
    "import sys\n",
    "print(sys.executable)\n",
    "\n",
    "%matplotlib inline"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2",
   "metadata": {},
   "source": [
    "## Parameters\n",
    "Set which waves to base the analysis on, what the minimum age must be to be considered and above what threshold a column is excluded due to missing data."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3",
   "metadata": {
    "tags": [
     "parameters"
    ]
   },
   "outputs": [],
   "source": [
    "s,m,e = 'a','q','s' # select which waves to base analysis on\n",
    "min_start_age = 25 # the minimum age people must as of the starting wave\n",
    "missing_threshold = 0.90\n",
    "correlation_threshold = 0.90\n",
    "redundant_threshold=0.9\n",
    "test = False\n",
    "release = \"restricted\" # \"restricted\""
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4",
   "metadata": {},
   "source": [
    "## Data Structure\n",
    "\n",
    "- Part1 contains the combined data from all questionairs asked in a given wave. Each wave is a separate file (eg a s wave 1, be is wave 2, etc. "
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5",
   "metadata": {},
   "source": [
    "### Filter people who were already studying"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6",
   "metadata": {
    "tags": [
     "hide-input"
    ]
   },
   "outputs": [],
   "source": [
    "summary_study = ['aedqstdy','aedfts','acaeft','acaept','anlreast','abncsty','abnfsty']\n",
    "\n",
    "c11_study = [\n",
    " 'aedcqsl',\n",
    " 'aedcqsh',\n",
    " 'aedcqnq',\n",
    " 'aedcqtq',\n",
    " 'aedcqta',\n",
    " 'aedcqtc',\n",
    " 'aedcqc1',\n",
    " 'aedcqc2',\n",
    " 'aedcqc3',\n",
    " 'aedcqc4',\n",
    " 'aedcqcd',\n",
    " 'aedcqad',\n",
    " 'aedcqav',\n",
    " 'aedcqbd',\n",
    " 'aedcqhd',\n",
    " 'aedcqgd',\n",
    " 'aedcqms',\n",
    " 'aedcqdc',\n",
    " 'aedcqbc',\n",
    " 'aedcqsc',\n",
    " 'aedcqcc',\n",
    " 'aedcqgc',\n",
    " 'aedcqcn',\n",
    " 'aedcqdn',\n",
    " 'aedcqnei',\n",
    " 'aedcqna',\n",
    " 'aedcqos',\n",
    " 'aedcqdk',\n",
    "]\n",
    "\n",
    "dv_asced_study = [\n",
    " 'aedcq100',\n",
    " 'aedcq110',\n",
    " 'aedcq120',\n",
    " 'aedcq200',\n",
    " 'aedcq211',\n",
    " 'aedcq221',\n",
    " 'aedcq310',\n",
    " 'aedcq311',\n",
    " 'aedcq312',\n",
    " 'aedcq400',\n",
    " 'aedcq411',\n",
    " 'aedcq413',\n",
    " 'aedcq421',\n",
    " 'aedcq500',\n",
    " 'aedcq511',\n",
    " 'aedcq514',\n",
    " 'aedcq521',\n",
    " 'aedcq524',\n",
    " 'aedcq600',\n",
    " 'aedcq611',\n",
    " 'aedcqunk'\n",
    "]\n",
    "\n",
    "already_studying_cols = summary_study + c11_study + dv_asced_study\n",
    "\n",
    "def filter_participants(df1,min_start_age, already_studying_cols):\n",
    "    \"\"\"\n",
    "    Remove those already studying or below the minimum age in the initial wave.\n",
    "    \"\"\"\n",
    "    n0 = len(df1)\n",
    "    df = df1.loc[df1[f'{s}hgage'] >= min_start_age].copy()\n",
    "    print(f\"Dropping {n0-len(df)} participants below age {min_start_age}\")\n",
    "\n",
    "    # filter out those already studying\n",
    "\n",
    "    already_studying = df[already_studying_cols].sum(axis=1)\n",
    "\n",
    "    n0 = len(df)\n",
    "    df = df[already_studying < 1].copy()\n",
    "    print(f\"Dropping {n0-len(df)} participants already studying at period start\")\n",
    "    print(f\"Remaining participants:{len(df)}\")\n",
    "    return df\n",
    "\n",
    "\n",
    "# read the combined file for the starting wave\n",
    "sfx = release_suffix[release]\n",
    "df1, meta1 = pyreadstat.read_sav(f'data/part1/Combined {s}190{sfx}.sav')\n",
    "n0 = len(df1)\n",
    "print(f\"Number of people in initial wave {n0}\")\n",
    "with open(f'data/metadata_{release}.pkl','wb') as f:\n",
    "    pickle.dump(meta1,f)\n",
    "    \n",
    "df1 = filter_participants(df1,min_start_age, already_studying_cols)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7",
   "metadata": {},
   "source": [
    "### Compute treatment & outcomes\n",
    "#### Outcomes measures\n",
    "   - **hours worked** Largly missing. The following variables are perfectly correlated\n",
    "      - `ajbhru` 57% missing, *E1b Hours per week usually worked in all your jobs*\n",
    "      - `ajbhruc`,57% missing *DV: Hours per week usually worked in all jobs*\n",
    "   - **wages (not normalised by hours worked)**\n",
    "      - `awsfe` wages from all jobs last financial year with imputation from net\n",
    "      - `awsce` current weekly wages from all jobs with imputation from net\n",
    "      - Non-imputed versions of both of these exist (replace the final `e` with `g`) but have slightly more missing values\n",
    "      - `awsfhave` records if people have received income from salary/wages last financial year. We could also use `_esbrd` to tell if people should have a non-zero wage.\n",
    "      - wage variables have quite a lot of missing data (~33% missing `awsfe` and 30% missing both `awsfe` and `awsfhave`)\n",
    "      - There are versions of wage variables with imputation of missing data from based on responses from the participant in other waves and responses from similar participants. These are indicated by the suffix `i` (eg `awsfei`, `awscei`). These variables contain no missing data.\n",
    "   - **employment status (categorical outcome)**\n",
    "   - **mental health**\n",
    "   \n",
    "#### Treatment measures\n",
    "   - Treatment is based on a change in education qualification between 2001 and 2017\n",
    "   - There are a number of study related variables that are only recorded on a subset of the waves. \n",
    "   - _edq{XXX} variables are recorded every year and count the number of qualifications a person holds in each of a number of categories\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8",
   "metadata": {
    "tags": [
     "hide_input"
    ]
   },
   "outputs": [],
   "source": [
    "from treatment_outcomes import compute_treatment_vars, compute_outcomes\n",
    "treatments = compute_treatment_vars(df1, s, m, release)\n",
    "outcomes = compute_outcomes(df1, s, e, release)\n",
    "treatment_outcomes = pd.merge(treatments,outcomes,on='xwaveid',how='inner')\n",
    "treatment_outcomes['xwaveid'] = treatment_outcomes['xwaveid'].astype(int)\n",
    "print(\"Treatments:\",treatments.columns)\n",
    "print(\"Outcomes:\",outcomes.columns)\n",
    "print(\"Updated computation of treatment\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9",
   "metadata": {},
   "source": [
    "#### Confusion matrix for highest vs count based treatment variable"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "10",
   "metadata": {
    "tags": [
     "hide_input"
    ]
   },
   "outputs": [],
   "source": [
    "from reed import compute_confusion\n",
    "compute_confusion(treatments['redudl'],treatments['reduhl'],'dl','dh')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "11",
   "metadata": {},
   "source": [
    "## Extract basic variables\n",
    "Extract a data set corresponding to the original paper we are working to extend - based on the table below;\n",
    "![image.png](images/original_paper_table.png)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "12",
   "metadata": {
    "tags": [
     "hide-input"
    ]
   },
   "outputs": [],
   "source": [
    "from treatment_outcomes import simplify_employment\n",
    "from reed import regex_select\n",
    "\n",
    "def extract_basic_variables(df):\n",
    "    # age, sex, education in 2001, employment status in 2001\n",
    "    basic = df1[['xwaveid','ahgage','ahgsex','aedhigh1','aesdtl']].copy() \n",
    "\n",
    "    def simplify_education(v):\n",
    "        \"\"\"Simplify down to match categories in paper.\"\"\"\n",
    "        if v < 0 or v==10:\n",
    "            return np.nan # missing\n",
    "        if v < 3: #(above bachelors)\n",
    "            return 2\n",
    "        return v # < year 12:(9), year 12:(8), cert:(5), diploma/adv diploma:(4), bachelors/honours:(3)\n",
    "    \n",
    "    # simplify education & employment in line with baseline paper\n",
    "    basic['aesdtl']=basic['aesdtl'].apply(simplify_employment)\n",
    "    basic['aedhigh1'] = basic['aedhigh1'].apply(simplify_education)\n",
    "    \n",
    "    # bin age\n",
    "    basic['ahgage'] = pd.cut(basic['ahgage'],bins=[24,34,44,54,120])\n",
    "    \n",
    "    # dummy encode\n",
    "    basic = pd.get_dummies(basic,columns=['ahgage','ahgsex','aedhigh1','aesdtl'],drop_first=True)\n",
    "    \n",
    "    # add interactions between gender and other variables\n",
    "    age_edu_emp = regex_select(basic.columns,['^ahgage_','^aedhigh1_','^aesdtl_'])\n",
    "    basic = create_interaction_columns(basic,['ahgsex_2.0'],age_edu_emp)\n",
    "    basic['xwaveid'] = basic['xwaveid'].astype(int)\n",
    "    return basic\n",
    "\n",
    "if not test:\n",
    "    basic = extract_basic_variables(df1)\n",
    "    l0 = len(basic)\n",
    "    basic_with_outcomes = pd.merge(basic,treatment_outcomes,on='xwaveid',how='inner')\n",
    "    l1 = len(basic_with_outcomes)\n",
    "    print(f\"Dropped {l0-l1} individuals who are not present in waves {m} and {e} ({100*(l0-l1)/l0:.0f}%)\")\n",
    "    basic_with_outcomes.set_index('xwaveid',inplace=True)\n",
    "    basic_with_outcomes.to_csv(\"data/basic_variables.csv\",index=True)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "13",
   "metadata": {},
   "source": [
    "## Extract Full variable set\n",
    "Extract a 'kitchen sink' dataset with minimal filtering of features"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "14",
   "metadata": {},
   "source": [
    "### Filter out columns based on annotated spreadsheet\n",
    "Remove columns that have been manually marked as irrelevant or proxies to whether someone is already studying. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "15",
   "metadata": {
    "tags": [
     "hide-input"
    ]
   },
   "outputs": [],
   "source": [
    "def read_type_information():\n",
    "    headers = ['variable','vartype','format','label','long_label','varcat','relevance',\"0\"]\n",
    "    type_df = pd.read_csv(\"data/HILDAw1vardic.csv\",skiprows=4,index_col=None, names=headers)\n",
    "    type_df['relevance'] = type_df['relevance'].fillna(1).astype(int)\n",
    "    type_df.loc[type_df['label']=='ACAEPT','relevance'] = -1\n",
    "    return type_df\n",
    "\n",
    "def drop_irrelevant_columns_inplace(df, type_df):\n",
    "    irrelevant = list(type_df.loc[type_df['relevance']<1,'variable'])\n",
    "    irrelevant.remove('xwaveid')\n",
    "    df.drop(columns=irrelevant,inplace=True)\n",
    "    print(f\"Dropped {len(irrelevant)} irrelevant columns.\")\n",
    "    return irrelevant"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "16",
   "metadata": {},
   "source": [
    "### Fix types\n",
    "   - encode categorical values & strings as integers (ordinal rather than one-hot)\n",
    "   - transform dates into days past 01/01/1900"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "17",
   "metadata": {
    "tags": [
     "hide-input"
    ]
   },
   "outputs": [],
   "source": [
    "def fix_types_inplace(df1):\n",
    "    # Reformat some of the columns\n",
    "    dates = [\n",
    "        \"ahhhqivw\",\n",
    "        \"ahhcompi\",\n",
    "        \"ahhcompf\",\n",
    "        \"ahhcomps\",\n",
    "        \"ahhidate\",\n",
    "        \"ahgdob1\",\n",
    "        \"ahgdob2\",\n",
    "        \"ahgdob3\",\n",
    "        \"ahgdob4\",\n",
    "        \"ahgdob5\",\n",
    "        \"ahgdob6\",\n",
    "        \"ahgdob\",\n",
    "    ]\n",
    "    string = ['ahhtitle']\n",
    "    categorical = [\n",
    "        'acca1',\n",
    "        'acca2',\n",
    "        'ahhmgfxd',\n",
    "        'ahhmgmxd',\n",
    "        'ahhp1',\n",
    "        'ahhp2',\n",
    "        'ahhp3',\n",
    "        'ahhpgfxd',\n",
    "        'ahhpgmxd',\n",
    "        'ahhpno',\n",
    "        'xwaveid',\n",
    "        'ahhid',\n",
    "        'ahhpid',\n",
    "        'ahhpcode',\n",
    "        'ahharea',\n",
    "        'ahhcd96',\n",
    "        'ats1',\n",
    "        'ats2',\n",
    "        'ats3',\n",
    "        'ats4',\n",
    "        'ats5',\n",
    "        'ats6',\n",
    "        'ahhdw',\n",
    "        'acsid1',\n",
    "        'acsid2',\n",
    "        'acsid3',\n",
    "        'acsid4',\n",
    "        'acsid5',\n",
    "        'acsid6',\n",
    "        'achid1',\n",
    "        'achid2',\n",
    "        'achid3',\n",
    "        'achid4',\n",
    "        'achid5',\n",
    "        'achid6',\n",
    "        'acpid1',\n",
    "        'acpid2',\n",
    "        'acpid3',\n",
    "        'acpid4',\n",
    "    ]\n",
    "\n",
    "    for c in categorical:\n",
    "        if c in df1.columns:\n",
    "            df1[c] = df1[c].replace(\"--\", \"\")\n",
    "            df1[c] = pd.to_numeric(df1[c])\n",
    "\n",
    "    # turn into days past epoch\n",
    "    basedate = pd.to_datetime('01/01/1900',format='%d/%m/%Y')    \n",
    "    for c in dates:\n",
    "        if c in df1.columns:\n",
    "            df1[c] = (pd.to_datetime(df1[c],format='%d/%m/%Y',errors='coerce')-basedate).dt.days \n",
    "\n",
    "    for c in string:\n",
    "        df1[c] = df1[c].astype('category').cat.codes"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "18",
   "metadata": {},
   "outputs": [],
   "source": [
    "def redundancy_with_nan(series, nbins):\n",
    "    \"\"\"\n",
    "    Compute the redundancy of a series after binning.\n",
    "    \n",
    "    Redundancy, in [0,1], measures how close the entropy of the distribution is to the maximum\n",
    "    entropy given the number of bins. Redundancy is minimized (=0) when the data is uniformly \n",
    "    distributed over the bins and maximised (=1) when all the data is in a single bin. \n",
    "    \"\"\"\n",
    "    if series.nunique() > nbins:\n",
    "        series = pd.cut(series, nbins, labels=False)\n",
    "            \n",
    "    counts = series.value_counts(dropna=False).values\n",
    "    if len(counts) == 1:\n",
    "        return 1\n",
    "    p = counts/counts.sum()\n",
    "    entropy = (-p*np.log(p)).sum()\n",
    "    redundancy = 1 - entropy/np.log(len(counts))\n",
    "    return redundancy\n",
    "\n",
    "def drop_redundant_columns_inplace(df,nbins, threshold):\n",
    "    r = np.zeros(len(df.columns))\n",
    "    for i,c in enumerate(df.columns):\n",
    "        r[i] = redundancy_with_nan(df[c],nbins)\n",
    "    \n",
    "    exclude = r > threshold\n",
    "    redundant = list(df.columns[exclude])\n",
    "    \n",
    "    r_values = dict(zip(df.columns[~exclude],r[~exclude]))\n",
    "    \n",
    "    df.drop(columns=redundant,inplace=True)\n",
    "    print(f\"Dropped {len(redundant)} columns with high redundancy/low entropy\")\n",
    "    return redundant, r_values"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "19",
   "metadata": {},
   "outputs": [],
   "source": [
    "def compute_correlations(df, fillval = 0):\n",
    "    \"\"\"Compute the correlations between each pair of variables and return as a DataFrame in long form.\"\"\"\n",
    "    c = df.fillna(fillval).corr()\n",
    "    c1 = []\n",
    "    c2 = []\n",
    "    value = []\n",
    "    for i in range(c.shape[0]):\n",
    "        for j in range(c.shape[0]):\n",
    "            if i > j:\n",
    "                value.append(c.iloc[i, j])\n",
    "                c1.append(c.index[i])\n",
    "                c2.append(c.columns[j])\n",
    "    c = pd.DataFrame({'c1': c1, 'c2': c2, \"correlation\": value})\n",
    "    c['abs'] = c['correlation'].abs()\n",
    "    c.sort_values(['abs'],ascending=False,inplace=True)\n",
    "    return c"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "20",
   "metadata": {},
   "outputs": [],
   "source": [
    "from collections import defaultdict\n",
    "\n",
    "def merge_correlated_pairs(df,r_vals, threshold, fillval=0):\n",
    "    \"\"\"\n",
    "    Merges pairs of variables with a correlation coefficient above the threshold.\n",
    "    \n",
    "    Returns\n",
    "    -------\n",
    "    merged: list[str]\n",
    "        A list of all the column names that were merged into other columns.\n",
    "        \n",
    "    merges: {str:[str]} \n",
    "        A dict from column name to all the columns merged with that column.\n",
    "    \"\"\"\n",
    "\n",
    "    cs = compute_correlations(df, fillval)\n",
    "    \n",
    "    column_indicies = [\n",
    "        cs.columns.get_loc('c1'),\n",
    "        cs.columns.get_loc('c2'),\n",
    "        cs.columns.get_loc('abs')\n",
    "    ]\n",
    "\n",
    "    row = cs.iloc[0,column_indicies]\n",
    "    c1, c2, t = row\n",
    "    merges = defaultdict(list)\n",
    "    while t >= threshold:\n",
    "\n",
    "        # merge (keep lowest redundancy)\n",
    "        r1, r2 = r_vals[c1], r_vals[c2]\n",
    "        if r1 <= r2:\n",
    "            best, other = c1, c2\n",
    "        else:\n",
    "            best, other = c2, c1\n",
    "\n",
    "        merges[best].append(other)\n",
    "\n",
    "        # delete all rows involving merged in variable\n",
    "        drop_index = cs.index[(cs['c1']==other)|(cs['c2']==other)]\n",
    "\n",
    "        cs.drop(index=drop_index,inplace=True)\n",
    "        row = cs.iloc[0,column_indicies]\n",
    "        c1, c2, t = row\n",
    "    \n",
    "    merged = []\n",
    "    for v in merges.values():\n",
    "        merged.extend(v)\n",
    "        \n",
    "    df.drop(columns=merged,inplace=True)\n",
    "    return merged, merges"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "21",
   "metadata": {},
   "source": [
    "### Automatic (very basic) column filtering\n",
    "   - drop columns where the proportion of data missing is above the maximum threshold\n",
    "   - drop columns that are constant (zero variance)\n",
    "   - drop columns that are very tightly correlated (based on correlation threshold) with another column that contains less missing data. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "22",
   "metadata": {
    "tags": [
     "hide-input"
    ]
   },
   "outputs": [],
   "source": [
    "def filter_raw_data(df1, missing_threshold=0.99, correlation_threshold=0.99):\n",
    "    columns_dropped = {} # keep track of why each column was dropped\n",
    "    \n",
    "    type_df = read_type_information()\n",
    "    irrelevant = drop_irrelevant_columns_inplace(df1,type_df)\n",
    "    add_list_to_dict(irrelevant,columns_dropped,'invalid/irrelevant')\n",
    "\n",
    "    fix_types_inplace(df1)\n",
    "    \n",
    "    constant = drop_constant_columns(df1)\n",
    "    add_list_to_dict(constant,columns_dropped,'constant')\n",
    "    \n",
    "    if redundant_threshold < 1:\n",
    "        redundant, r_vals = drop_redundant_columns_inplace(df1, 100, redundant_threshold)\n",
    "        add_list_to_dict(redundant,columns_dropped,'high-redundancy')\n",
    "    \n",
    "    if missing_threshold < 1:\n",
    "        mostly_missing = drop_mostly_missing_columns(df1, threshold = missing_threshold)\n",
    "        add_list_to_dict(mostly_missing, columns_dropped, 'mostly-missing')\n",
    "    \n",
    "    if correlation_threshold < 1:\n",
    "        dropped, merges = merge_correlated_pairs(df1, r_vals, correlation_threshold,fillval=0)\n",
    "        add_list_to_dict(dropped, columns_dropped, 'merged')\n",
    "    \n",
    "    \n",
    "    print(\"Processed data, with shape:\",df1.shape)\n",
    "    return df1, columns_dropped, r_vals\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "23",
   "metadata": {},
   "source": [
    "### Save the data to file\n",
    "Save the data to file for subsequent model fitting"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "24",
   "metadata": {
    "tags": [
     "hide_input"
    ]
   },
   "outputs": [],
   "source": [
    "import os\n",
    "def write_data(X, treatment_outcomes,filename):\n",
    "    filepath = os.path.join(\"data\",filename)\n",
    "    l0 = len(X)\n",
    "    df = pd.merge(X, treatment_outcomes, on=['xwaveid'],how='inner')\n",
    "    l1 = len(df)\n",
    "    print(f\"Dropped {l0-l1} individuals who are not present in waves {m} and {e} ({100*(l0-l1)/l0:.0f}%)\")\n",
    "    df.set_index('xwaveid',inplace=True)\n",
    "    df.to_csv(filepath,index=True)\n",
    "    print(f\"Written data of shape {df.shape} to:\",filepath)\n",
    "    assert (basic_with_outcomes.index == df.index).all(), \"index should be the same across datasets\"\n",
    "\n",
    "if not test:\n",
    "    X, columns_dropped, r_vals = filter_raw_data(df1.copy(), missing_threshold=missing_threshold,correlation_threshold=correlation_threshold)\n",
    "    write_data(X, treatment_outcomes, f\"all_vars_{release}.csv\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "25",
   "metadata": {},
   "source": [
    "## Feature selection using wage 5 years after initial wave as labels\n",
    "\n",
    "This is based on the premise that the features that are predictive of the change in wage due to returning to education are likely to be the same set of features that are important for predicting wage more generally. This allows us to do supervised feature selection without worrying so much about over-fitting to the data, as we are not using the final labels we are training against. \n",
    "\n",
    "Note: we can't use initial wage as the target because then we end up selecting features that are simply proxies to the target."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "26",
   "metadata": {},
   "outputs": [],
   "source": [
    "from cinspect.dimension import effective_rank\n",
    "from sklearn.feature_selection import SelectKBest, mutual_info_regression, f_regression\n",
    "from sklearn.preprocessing import StandardScaler\n",
    "import string"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "27",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Make a dataset for feature selection that uses the value of the target feature 5 years after the initial wave as the target\n",
    "\n",
    "def create_dataset_to_predict_outcome_five_years_after_initial_wave(target):\n",
    "    indx = string.ascii_lowercase.index(s)+5\n",
    "    post_start = string.ascii_lowercase[indx]\n",
    "   \n",
    "    df0, meta1 = pyreadstat.read_sav(f'data/part1/Combined {s}190{sfx}.sav') \n",
    "    outcomes_post = compute_outcomes(df0, s, post_start, release)\n",
    "    outcomes_post['xwaveid'] = outcomes_post['xwaveid'].astype(int)\n",
    "    outcomes_post.set_index('xwaveid',inplace=True)\n",
    "    X, columns_dropped, r_vals = filter_raw_data(df0, missing_threshold=0.9,correlation_threshold=0.9)\n",
    "    X['xwaveid'] = X['xwaveid'].astype(int)\n",
    "    X.set_index('xwaveid',inplace=True)\n",
    "    o = outcomes_post.join(X)\n",
    "    features = df0.columns\n",
    "    y = o[target].values\n",
    "    valid_rows = ~np.isnan(y)\n",
    "\n",
    "    y = y[valid_rows]\n",
    "    X = o.loc[valid_rows,features]\n",
    "    Xs = StandardScaler().fit_transform(X.fillna(0))\n",
    "    Xs = pd.DataFrame(Xs, columns=features)\n",
    "    return Xs, y\n",
    "\n",
    "if not test:\n",
    "    target = 'y_wscei'\n",
    "    Xs, ys = create_dataset_to_predict_outcome_five_years_after_initial_wave(target)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "28",
   "metadata": {},
   "source": [
    "### Univariate features selectiong: Select k-best\n",
    "\n",
    "Takes no account of correlation between features, results in a matrix with a low effective rank (compared with it's actual rank). "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "29",
   "metadata": {},
   "outputs": [],
   "source": [
    "if not test:\n",
    "    selector = SelectKBest(f_regression, k=10)\n",
    "    selector.fit(Xs,ys)\n",
    "    effective_rank(Xs[selector.get_feature_names_out()])"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "30",
   "metadata": {},
   "source": [
    "### Multivariate-feature selection: Lasso"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "31",
   "metadata": {},
   "outputs": [],
   "source": [
    "# see how many features we get as a function of regularisation strength. \n",
    "# we could use prediction accuracy as a rough estimate of how many features to use\n",
    "# note that we are not removing people currently studying, so the model may select features relating to current study\n",
    "# these will be dropped before being written as they are constant. \n",
    "from sklearn.linear_model import Lasso, LassoCV\n",
    "from sklearn.feature_selection import SelectFromModel\n",
    "if not test:\n",
    "    alphas = np.logspace(-1,2,40)  \n",
    "    n_selected = []\n",
    "    for alpha in alphas:\n",
    "        ls = Lasso(alpha=alpha)\n",
    "        selector = SelectFromModel(ls)\n",
    "        selector.fit(Xs,ys)\n",
    "        n_selected.append(len(selector.get_feature_names_out()))\n",
    "\n",
    "    plt.semilogx(alphas, n_selected)\n",
    "    plt.xlabel(\"alpha\")\n",
    "    plt.ylabel(\"number of features selected\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "32",
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_alpha_for_num_features(num_features, n_selected, alphas):\n",
    "    n_selected = np.array(n_selected)\n",
    "    indx = np.nanargmin(np.abs(n_selected - num_features))\n",
    "    return alphas[indx], n_selected[indx]\n",
    "\n",
    "def select_features(num_features, n_selected, alphas, Xs, ys):\n",
    "    print(f\"Generating dataset with approximately {num_features} features\" )\n",
    "    alpha, _ = get_alpha_for_num_features(num_features,n_selected, alphas)\n",
    "    ls = Lasso(alpha=alpha)\n",
    "    selector = SelectFromModel(ls)\n",
    "    selector.fit(Xs,ys)\n",
    "    f_selected = selector.get_feature_names_out()\n",
    "    print(\"number of features selected:\",len(f_selected))\n",
    "    print(\"effective rank on selection matrix:\",effective_rank(Xs[f_selected]))\n",
    "\n",
    "    print(\"\\n\")\n",
    "    return f_selected\n",
    "\n",
    "def write_selected_featureset(X, columns_dropped, features, tag, release):\n",
    "    selection = ['xwaveid']\n",
    "    missing = []\n",
    "    for f in features:\n",
    "        if f in X.columns:\n",
    "            selection.append(f)\n",
    "        else:\n",
    "            reason = columns_dropped.get(f,\"unkown\")\n",
    "            print(f\"Column {f} not present in X, reason:{reason}\")\n",
    "    \n",
    "    write_data(X[selection], treatment_outcomes, f\"all_lasso_selected_{tag}_{release}.csv\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "33",
   "metadata": {},
   "outputs": [],
   "source": [
    "if not test:\n",
    "    f10 = select_features(10, n_selected, alphas, Xs, ys)\n",
    "    f20 = select_features(20, n_selected, alphas, Xs, ys)\n",
    "    f50 = select_features(50, n_selected, alphas, Xs, ys)\n",
    "    f100 = select_features(100, n_selected, alphas, Xs, ys)\n",
    "\n",
    "    X, columns_dropped, r_vals = filter_raw_data(df1.copy(), missing_threshold=0.99,correlation_threshold=1)\n",
    "\n",
    "    write_selected_featureset(X, columns_dropped, f10, 10, release)\n",
    "    write_selected_featureset(X, columns_dropped, f20, 20, release)\n",
    "    write_selected_featureset(X, columns_dropped, f50, 50, release)\n",
    "    write_selected_featureset(X, columns_dropped, f100, 100, release)\n",
    "    for f in f20:\n",
    "        print(f, meta1.column_names_to_labels.get(f))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "34",
   "metadata": {},
   "source": [
    "## Unsupervised feature selection\n",
    "\n",
    "Greedy selection of features to optimise effective rank."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "35",
   "metadata": {},
   "outputs": [],
   "source": [
    "# from cinspect.dimension import greedy_feature_selection, effective_rank\n",
    "\n",
    "# features_to_select = 100\n",
    "# Xs = StandardScaler().fit_transform(X.fillna(0))\n",
    "# selected, vals = greedy_feature_selection(\n",
    "#     Xs, effective_rank, \n",
    "#     initial_col=4, num_to_select=features_to_select\n",
    "# )\n",
    "\n",
    "# selected_cols = [X.columns[i] for i in selected]\n",
    "# X_su = X[selected_cols]\n",
    "# print(\"number of features selected:\",len(selected_cols))\n",
    "# print(\"effective rank:\",effective_rank(Xs[selected]))\n",
    "# write_data(X_su.reset_index(), treatment_outcomes, f\"all_unsupervised_selected_{release}.csv\")"
   ]
  }
 ],
 "metadata": {
  "celltoolbar": "Tags",
  "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.7"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}