add example ukb RAP convert notebook

Author: tf2

data/ukb_simulated_data/example_ukb_rap_convert.ipynb

@@ -0,0 +1,235 @@
+{
+ "cells": [
+  {
+   "cell_type": "markdown",
+   "id": "c6606cb2",
+   "metadata": {},
+   "source": [
+    "# General example for converting UK Biobank data into the research environemnt (RAP) to DELPHI format\n"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "ba4bfe9e",
+   "metadata": {},
+   "source": [
+    "Notes:\n",
+    " - This is setup to run as a notebook in a spark server job\n",
+    " - it needs access to a dataset record (which you may need to explicitly specify for your project)\n",
+    " - and also a cohort which is here refered to as \"full_cohort\" this may differ from your project\n",
+    " - the token labels to ukb field ids should be modified to suit your needs\n",
+    " - your should change the token ids to be integer and create a token_id to field_id (or disease icd10 name)\n",
+    " - this output a file with all individuals included - you will want to split this into \"train.bin\" and \"val.bin\" "
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "3c85eb2e",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "import dxdata\n",
+    "import pandas as pd\n",
+    "import numpy as np\n",
+    "from tqdm import tqdm\n",
+    "from pyspark.sql.functions import lit, udf, col\n",
+    "from pyspark.sql.types import DoubleType\n",
+    "from functools import reduce\n",
+    "import os\n",
+    "\n",
+    "def get_first_occ_fields(main_entity):\n",
+    "    fo_fields = []\n",
+    "    for field in main_entity.fields:\n",
+    "        parts = field.name.split(\"_\")\n",
+    "        if len(str(parts[0])) > 3:\n",
+    "            field_num = int(parts[0][1:])\n",
+    "            if (field_num >= 130000 and field_num <= 132604):\n",
+    "                if field.title.startswith(\"Date\"):\n",
+    "                    fo_fields.append(field)\n",
+    "    return fo_fields\n",
+    "\n",
+    "def compute_age_from_eid_and_event(eid, event_date):\n",
+    "    dob = dob_lookup.get(eid)\n",
+    "    if dob is None or event_date is None:\n",
+    "        return None\n",
+    "    try:\n",
+    "        return (pd.to_datetime(event_date) - dob).days / 365.25\n",
+    "    except Exception:\n",
+    "        return None\n",
+    "\n",
+    "# Initialize dxdata engine\n",
+    "engine = dxdata.connect(dialect=\"hive+pyspark\")\n",
+    "\n",
+    "project = os.getenv('DX_PROJECT_CONTEXT_ID')\n",
+    "record = os.popen(\"dx find data --type Dataset --delimiter ',' | awk -F ',' '{print $5}'\").read().rstrip()\n",
+    "# find what is presumed to be the relevant dataset record\n",
+    "record = record.split('\\n')[0]\n",
+    "\n",
+    "DATASET_ID = project + \":\" + record\n",
+    "dataset = dxdata.load_dataset(id=DATASET_ID)\n",
+    "\n",
+    "# we retrieve the priamry entity from the dataset\n",
+    "main_entity = dataset.primary_entity\n",
+    "\n",
+    "# use cohort - change to whichever name:path you have for this object\n",
+    "cohort = dxdata.load_cohort(folder=\"/\", name=\"full_cohort\")\n",
+    "cohort_eids_df = engine.execute(cohort.sql)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "361f018d-fe1c-40ec-b2f4-d43f7c9887ca",
+   "metadata": {
+    "tags": []
+   },
+   "outputs": [],
+   "source": [
+    "# hard coded sex dob and basic demographic data\n",
+    "eid_f = main_entity.find_field(name=\"eid\")\n",
+    "sex_f = main_entity.find_field(title=\"Sex\")\n",
+    "year_f = main_entity.find_field(title=\"Year of birth\")\n",
+    "month_f = main_entity.find_field(title=\"Month of birth\")\n",
+    "death_f = dataset['death'].find_field(title=\"Date of death\")\n",
+    "assessment_f = main_entity[\"p53_i0\"]\n",
+    "bmi_f = main_entity[\"p21001_i0\"]\n",
+    "smoking_f = main_entity[\"p1239_i0\"]\n",
+    "alcohol_f = main_entity[\"p1558_i0\"]\n"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 5,
+   "id": "004328b8-abfe-4e63-b704-87efbe7138cd",
+   "metadata": {
+    "tags": []
+   },
+   "outputs": [],
+   "source": [
+    "# collect the cancer code enteries\n",
+    "cancer_codes = {}\n",
+    "cancer_codes['type'] = []\n",
+    "cancer_codes['date'] = []\n",
+    "for i in range(22):\n",
+    "    cancer_codes['type'].append(main_entity.find_field(name=\"p40006_i\" + str(i)))\n",
+    "    cancer_codes['date'].append(main_entity.find_field(name=\"p40005_i\" + str(i)))\n"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "0381027b-de51-494e-aa12-d3a0c126d35f",
+   "metadata": {
+    "tags": []
+   },
+   "outputs": [],
+   "source": [
+    "# Deal with first occrances and demographic data - this takes a little while\n",
+    "fo_fields = get_first_occ_fields(main_entity)\n",
+    "fields_to_get = [eid_f, sex_f, year_f, month_f, assessment_f, bmi_f, smoking_f, alcohol_f] + cancer_codes['type'] + cancer_codes['date'] + fo_fields + [death_f]\n",
+    "\n",
+    "df = main_entity.retrieve_fields(fields=fields_to_get, filter_sql=cohort.sql, engine=engine)\n",
+    "df1 = df.select(\"eid\", \"p31\",\"p34\",\"p52\",\"p53_i0\",\"p21001_i0\",\"p1239_i0\",\"p1558_i0\").toPandas()\n",
+    "\n",
+    "dobf1  = df1[['p34', 'p52']]\n",
+    "dobf1.columns = [\"YEAR\", \"MONTH\"]\n",
+    "df1['dob'] = pd.to_datetime(dobf1.assign(DAY=1))\n",
+    "df1['bmi_status'] = np.where(df1['p21001_i0']>28,5,np.where(df1['p21001_i0']>22,4,3))\n",
+    "df1['smoking_status'] = np.where(df1['p1239_i0']==1,8,np.where(df1['p1239_i0']==2,7,6))\n",
+    "df1['alcohol_status'] = np.where(df1['p1558_i0']==1,11,np.where(df1['p1558_i0'] < 4,10,9))"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "343290db-0f6c-4717-bc25-b821005c0c31",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# Prepare a pandas dictionary for fast eid to dob lookup\n",
+    "dob_lookup = df1.set_index('eid')['dob'].to_dict()\n",
+    "age_event_udf = udf(compute_age_from_eid_and_event, DoubleType())"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "d0a7b99c-7dd9-43c8-9766-b10256282117",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# Remove all NULL enteries for each ICD10 code seperately and combine into an overall spark table \n",
+    "\n",
+    "# deal with the tokens and dates \n",
+    "d_all = df.select(\"eid\", cancer_codes['date'][0].name, cancer_codes['type'][0].name).where(df[cancer_codes['date'][0].name].isNotNull())\n",
+    "d_all = d_all.withColumnRenamed(cancer_codes['date'][i].name, \"date\")\n",
+    "d_all = d_all.withColumnRenamed(cancer_codes['type'][i].name, \"token\")\n",
+    "d_all = d_all.withColumn(\"age\", age_event_udf(col(\"eid\"), col(\"date\")))\n",
+    "    \n",
+    "for i in tqdm(1, len(cancer_codes['date'])):\n",
+    "    cf1 = df.select(\"eid\", cancer_codes['date'][i].name, cancer_codes['type'][i].name).where(df[cancer_codes['date'][i].name].isNotNull())\n",
+    "    cf1 = cf1.withColumnRenamed(cancer_codes['date'][i].name, \"date\")\n",
+    "    cf1 = cf1.withColumnRenamed(cancer_codes['type'][i].name, \"token\")\n",
+    "    cf1 = cf1.withColumn(\"age\", age_event_udf(col(\"eid\"), col(\"date\")))\n",
+    "    d_all = d_all.union(cf1)\n"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "edeb7116-d937-4e8a-90dd-fff159e64a07",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# Deal with all first occurances - this takes a long time\n",
+    "for i in tqdm(range(0,len(fo_fields))):\n",
+    "    f = fo_fields[i]\n",
+    "    d = df.select(['eid', f.name]).where(df[f.name].isNotNull())\n",
+    "    d1 = d.withColumn(\"token\", lit(f.name))\n",
+    "    d1 = d1.withColumnRenamed(f.name, \"date\")\n",
+    "    d1 = d1.withColumn(\"age\", age_event_udf(col(\"eid\"), col(\"date\")))\n",
+    "    d_all = d_all.union(d1)\n"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "0770c44a-7585-42b3-8f48-cadba006c54f",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# Format, sort and write out to a file\n",
+    "df_all = d_all.select(\"eid\", \"age\", \"token\").toPandas()\n",
+    "df_all['age'] = df_all['age'] * 365.25\n",
+    "data = np.array(df_all).squeeze()\n",
+    "data[:,0] = data[:,0].astype(np.uint32)\n",
+    "data[:,1] = data[:,1].astype(np.uint32)\n",
+    "data = data[np.lexsort((data[:,1], data[:,0]))]\n",
+    "\n",
+    "data.tofile('all_records.bin')\n"
+   ]
+  }
+ ],
+ "metadata": {
+  "kernelspec": {
+   "display_name": "Python 3 (ipykernel)",
+   "language": "python",
+   "name": "python3"
+  },
+  "language_info": {
+   "codemirror_mode": {
+    "name": "ipython",
+    "version": 3
+   },
+   "file_extension": ".py",
+   "mimetype": "text/x-python",
+   "name": "python",
+   "nbconvert_exporter": "python",
+   "pygments_lexer": "ipython3",
+   "version": "3.12.9"
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 5
+}