scripts/data_cleaning_and_splitting.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "632e9cd2",
   "metadata": {},
   "outputs": [],
   "source": [
    "import pandas as pd\n",
    "import tqdm\n",
    "import os\n",
    "import numpy as np\n",
    "import sys\n",
    "import copy\n",
    "\n",
    "from collections import Counter\n",
    "\n",
    "from sklearn.model_selection import train_test_split\n",
    "\n",
    "from pandarallel import pandarallel\n",
    "pandarallel.initialize(progress_bar=True)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7bb9ac80-576c-48e9-8455-31fc574694a7",
   "metadata": {},
   "source": [
    "# Input here\n",
    "\n",
    "The path to the folder where Swedish MS REGistry (SMSREG) is kept.  "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0d960fd0",
   "metadata": {},
   "outputs": [],
   "source": [
    "filepath = \"\""
   ]
  },
  {
   "cell_type": "markdown",
   "id": "375a3223",
   "metadata": {},
   "source": [
    "# Demographic data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4c0145c6-a71a-45c6-9d51-e81e068558af",
   "metadata": {},
   "outputs": [],
   "source": [
    "# read base file and visit file\n",
    "basdata_df = pd.read_csv(filepath + \"basdata.csv\")\n",
    "besoksdata_df = pd.read_csv(filepath + \"besoksdata.csv\")\n",
    "\n",
    "# Checking patient and visit counts\n",
    "print (\"Total number of patients in the data =\",len(basdata_df[\"patient code\"].unique()))\n",
    "print (\"Total number of visits in the data =\",len(besoksdata_df[besoksdata_df[\"patient code\"].isin(basdata_df[\"patient code\"].unique())]))\n",
    "print (\"Total number of PPMS patients\",len(basdata_df[basdata_df[\"progress\"] == \"PP - Primary Progressive\"]))\n",
    "print (\"Total number of RRMS and SPMS patients\",len(basdata_df[~(basdata_df[\"progress\"] == \"PP - Primary Progressive\")]))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f3aacc66",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Get progress and to sp from unique entry\n",
    "# get pcode, debut date, diagnosis date, to sp date, date of birth, current progress, sex, visit date\n",
    "# Convert them to required data format\n",
    "PatientCode_Progress_df = basdata_df[basdata_df[\"patient code\"].notna()][[\"patient code\",\"progress\",\"to sp\",\"birthday\",\"sex\",\"debut date\",\"diagnosis date\"]]\n",
    "PatientCode_Progress_df[\"debut_date\"] =  pd.to_datetime(PatientCode_Progress_df['debut date'], format='%Y-%m-%d')\n",
    "PatientCode_Progress_df[\"diagnosis_date\"] =  pd.to_datetime(PatientCode_Progress_df['diagnosis date'], format='%Y-%m-%d')\n",
    "PatientCode_Progress_df[\"to_sp\"] =  pd.to_datetime(PatientCode_Progress_df['to sp'], format='%Y')\n",
    "PatientCode_Progress_df[\"birthday\"] =  pd.to_datetime(PatientCode_Progress_df['birthday'], format='%Y-%m-%d')\n",
    "PatientCode_Progress_df = PatientCode_Progress_df[PatientCode_Progress_df[\"progress\"].notna()]\n",
    "PatientCode_Progress_df = PatientCode_Progress_df[PatientCode_Progress_df[\"sex\"].notna()]\n",
    "PatientCode_Progress_df = PatientCode_Progress_df.drop([\"diagnosis date\"],axis=1)\n",
    "PatientCode_Progress_df = PatientCode_Progress_df.drop([\"debut date\"],axis=1)\n",
    "PatientCode_Progress_df = PatientCode_Progress_df.drop([\"to sp\"],axis=1)\n",
    "\n",
    "PatientCode_VisitDate_EDSS_Age_df = besoksdata_df[[\"patient code\",\"date\",\"edss value/score\"]].merge(basdata_df[[\"patient code\",\"birthday\"]], on='patient code', how='left')\n",
    "PatientCode_VisitDate_EDSS_Age_df[\"visit_date\"] =  pd.to_datetime(PatientCode_VisitDate_EDSS_Age_df['date'], format='%Y-%m-%d')\n",
    "PatientCode_VisitDate_EDSS_Age_df = PatientCode_VisitDate_EDSS_Age_df.drop([\"date\"],axis=1)\n",
    "PatientCode_VisitDate_EDSS_Age_df[\"birthday\"] =  pd.to_datetime(PatientCode_VisitDate_EDSS_Age_df['birthday'], format='%Y-%m-%d')\n",
    "PatientCode_VisitDate_EDSS_Age_df[\"age_at_visit\"] = np.round((PatientCode_VisitDate_EDSS_Age_df[\"visit_date\"]-PatientCode_VisitDate_EDSS_Age_df[\"birthday\"]) / np.timedelta64(1, 'Y'),3)\n",
    "PatientCode_VisitDate_EDSS_Age_df = PatientCode_VisitDate_EDSS_Age_df[[\"patient code\",\"visit_date\",\"edss value/score\",\"age_at_visit\"]]\n",
    "PatientCode_VisitDate_EDSS_Age_df = PatientCode_VisitDate_EDSS_Age_df.dropna()\n",
    "\n",
    "# Duplicate visits to be removed \n",
    "PatientCode_VisitDate_EDSS_Age_df  = PatientCode_VisitDate_EDSS_Age_df.drop_duplicates(subset=['patient code', 'visit_date'],keep=\"last\")\n",
    "demographic_df = PatientCode_VisitDate_EDSS_Age_df.merge(PatientCode_Progress_df, on='patient code', how='left')\n",
    "\n",
    "demographic_df[\"age_at_debut\"] = np.round((demographic_df[\"debut_date\"]-demographic_df[\"birthday\"]) / np.timedelta64(1, 'Y'),3)\n",
    "demographic_df[\"age_at_diagnosis\"] = np.round((demographic_df[\"diagnosis_date\"]-demographic_df[\"birthday\"]) / np.timedelta64(1, 'Y'),3)\n",
    "demographic_df[\"age_at_to_sp\"] = np.round((demographic_df[\"to_sp\"]-demographic_df[\"birthday\"]) / np.timedelta64(1, 'Y'),3)\n",
    "demographic_df = demographic_df.dropna(subset=['patient code', 'visit_date', 'edss value/score', 'age_at_visit', 'progress', 'birthday', 'sex'])\n",
    "demographic_df= demographic_df.sort_values(by=[\"patient code\",'visit_date'])\n",
    "\n",
    "demographic_df = demographic_df[~(demographic_df[\"debut_date\"].isna() & demographic_df[\"diagnosis_date\"].isna())] # Remove if both debut and diagnosis dates are missing\n",
    "print (\"Number of person in demographic data = \",len(set(demographic_df[\"patient code\"].to_list())))\n",
    "print (\"Total number of entries = \",len(demographic_df[\"patient code\"].to_list()))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f06b3fc9",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Assigning RRMS or SPMS or PPMS for each visits\n",
    "demographic_df[\"days_diff\"] = (demographic_df[\"visit_date\"] - demographic_df[\"to_sp\"])/ np.timedelta64(1, 'D')\n",
    "demographic_df.loc[demographic_df[\"days_diff\"] < 0.1,\"progress during visit\"] = \"RR\"\n",
    "demographic_df.loc[demographic_df[\"days_diff\"] > 0.1,\"progress during visit\"] = \"SP\"\n",
    "demographic_df.loc[demographic_df[\"progress\"] == \"PP - Primary Progressive\",\"progress during visit\"] = \"PP\"\n",
    "demographic_df.loc[demographic_df[\"progress\"] == \"RR - Intermittent progressive MS\",\"progress during visit\"] = \"RR\"\n",
    "demographic_df = demographic_df.drop([\"days_diff\"],axis=1)\n",
    "\n",
    "# Printing statistics of the data\n",
    "print (\"Total number of datapoints in the df = \", len(demographic_df))\n",
    "demographic_df = demographic_df[demographic_df[\"progress during visit\"].notna()]\n",
    "pp_snapshots = demographic_df[(demographic_df[\"progress during visit\"] == \"PP\")]\n",
    "rr_snapshots = demographic_df[(demographic_df[\"progress during visit\"] == \"RR\")]\n",
    "sp_snapshots = demographic_df[demographic_df[\"progress during visit\"] == \"SP\"]\n",
    "print (\"Total number of datapoints in the df after removing no sp/pp/rr info= \", len(demographic_df))\n",
    "sp_patient_list = list(set(sp_snapshots[\"patient code\"].to_list()))\n",
    "rr_patient_list = list(set(rr_snapshots[\"patient code\"].to_list()))\n",
    "rr_patient_list = list(set(rr_patient_list) - set(sp_patient_list))\n",
    "pp_patient_list = list(set(pp_snapshots[\"patient code\"].to_list()))\n",
    "print (\"Total number of people in PP snapshots = \", len(set(pp_patient_list)))\n",
    "print (\"Total number of people in RR snapshots = \", len(set(rr_patient_list)))\n",
    "print (\"Total number of people in SP snapshots = \", len(set(sp_patient_list)))\n",
    "print (\"Total number of PP snapshots = \", len(demographic_df[demographic_df[\"patient code\"].isin(pp_patient_list)]))\n",
    "print (\"Total number of RR snapshots = \", len(demographic_df[demographic_df[\"patient code\"].isin(rr_patient_list)]))\n",
    "print (\"Total number of SP snapshots = \", len(demographic_df[demographic_df[\"patient code\"].isin(sp_patient_list)]))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2b7d91a0",
   "metadata": {},
   "source": [
    "# EQ5D"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6af1c1c2",
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "demographic_addinfo_df = copy.deepcopy(demographic_df)\n",
    "eq5d_df = pd.read_csv(filepath + \"eq5d.csv\")\n",
    "eq5d_df[\"date\"] =  pd.to_datetime(eq5d_df['date'], format='%Y-%m-%d',errors = \"coerce\") # Change to date format\n",
    "eq5d_df = eq5d_df[~eq5d_df[\"date\"].isna()] # Remove all the NaT\n",
    "\n",
    "eq5d_df = eq5d_df.dropna(subset=[\"date\"]) # Remove if date is not present\n",
    "eq5d_df= eq5d_df.sort_values(\"date\")\n",
    "\n",
    "def get_last_eq5d_using_date(visit_date,patient_code):\n",
    "    df = eq5d_df[eq5d_df[\"patient code\"]==patient_code][eq5d_df[eq5d_df[\"patient code\"]==patient_code][\"date\"] <= visit_date]\n",
    "    df= df.sort_values(\"date\")\n",
    "    if len(df) > 0:\n",
    "        last_eq5d_date = df[\"date\"].to_list()[-1]\n",
    "        eq5d_score = df[\"eq5d score\"].to_list()[-1]\n",
    "        return pd.Series([last_eq5d_date,eq5d_score])\n",
    "    return pd.Series([np.nan,np.nan])\n",
    "\n",
    "demographic_addinfo_df[[\"eq5d_date\",'eq5d_score']] = demographic_addinfo_df.parallel_apply(lambda x: get_last_eq5d_using_date(x[\"visit_date\"],x[\"patient code\"]),axis=1)\n",
    "#demographic_addinfo_df = demographic_addinfo_df[demographic_addinfo_df[\"sdmt days\"].notna()] # remove if sdmt== nan\n",
    "demographic_addinfo_df[\"age_at_eq5d\"] = np.round((demographic_addinfo_df[\"eq5d_date\"]-demographic_addinfo_df[\"birthday\"]) / np.timedelta64(1, 'Y'),3)\n",
    "demographic_addinfo_df.loc[demographic_addinfo_df[\"age_at_eq5d\"].notna() & demographic_addinfo_df[\"eq5d_score\"].isna(),\"age_at_eq5d\"] = np.nan # Remove when only visit is present and not score"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0d1ef087",
   "metadata": {},
   "source": [
    "# SDMT"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "51345fe0",
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "sdmt_df = pd.read_csv(filepath + \"sdmt.csv\")\n",
    "sdmt_df[\"date\"] =  pd.to_datetime(sdmt_df['date'], format='%Y-%m-%d',errors = \"coerce\")\n",
    "sdmt_df = sdmt_df[sdmt_df[\"date\"].notna() & sdmt_df[\"number passed\"].notna()]\n",
    "\n",
    "sdmt_df = sdmt_df.dropna(subset=[\"date\"]) # Remove if date is not present\n",
    "sdmt_df= sdmt_df.sort_values(\"date\")\n",
    "\n",
    "def get_last_sdmt_using_date(visit_date,patient_code):\n",
    "    df = sdmt_df[sdmt_df[\"patient code\"]==patient_code][sdmt_df[sdmt_df[\"patient code\"]==patient_code][\"date\"] <= visit_date]\n",
    "    df= df.sort_values(\"date\")\n",
    "    if len(df) > 0:\n",
    "        last_sdmt_date = df[\"date\"].to_list()[-1]\n",
    "        sdmt_days_passed = df[\"number passed\"].to_list()[-1]\n",
    "        \n",
    "        return pd.Series([last_sdmt_date,sdmt_days_passed])\n",
    "    return pd.Series([np.nan,np.nan])\n",
    "\n",
    "demographic_addinfo_df[[\"sdmt_date\",'sdmt_score']] = demographic_addinfo_df.parallel_apply(lambda x: get_last_sdmt_using_date(x[\"visit_date\"],x[\"patient code\"]),axis=1)\n",
    "#demographic_addinfo_df = demographic_addinfo_df[demographic_addinfo_df[\"sdmt days\"].notna()] # remove if sdmt== nan\n",
    "demographic_addinfo_df[\"age_at_sdmt\"] = np.round((demographic_addinfo_df[\"sdmt_date\"]-demographic_addinfo_df[\"birthday\"]) / np.timedelta64(1, 'Y'),3)\n",
    "demographic_addinfo_df.loc[demographic_addinfo_df[\"age_at_sdmt\"].notna() & demographic_addinfo_df[\"sdmt_score\"].isna(),\"age_at_sdmt\"] = np.nan"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3701dc36",
   "metadata": {},
   "source": [
    "# Treatments"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "bb919c95",
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_terapi_drug(patient_code,visit_date):\n",
    "    current_vist_date_drugs = []\n",
    "    for index,row in terapi_df[terapi_df[\"patient code\"]==patient_code].iterrows():\n",
    "        start_date = row[\"insert_date\"]\n",
    "        end_date = row[\"end_date\"]\n",
    "        if (start_date - visit_date).days < 0 and (end_date - visit_date).days >= 0:\n",
    "            drug_given = row[\"preparation\"]\n",
    "            current_vist_date_drugs.append(drug_given)\n",
    "        elif (start_date - visit_date).days < 0 and str(end_date - visit_date) == \"NaT\":\n",
    "            drug_given = row[\"preparation\"]\n",
    "            current_vist_date_drugs.append(drug_given)\n",
    "    if len(current_vist_date_drugs) == 1:\n",
    "        if current_vist_date_drugs[0] == \"Ingen behandling\":\n",
    "            return \"no treatment\"\n",
    "        return current_vist_date_drugs[0]\n",
    "    elif len(current_vist_date_drugs) == 0:\n",
    "        return \"no treatment\"\n",
    "    return current_vist_date_drugs\n",
    "\n",
    "\n",
    "terapi_df = pd.read_csv(filepath + \"terapi.csv\")\n",
    "terapi_df = terapi_df[[\"patient code\",\"insert\",\"preparation\",\"exposed\"]]\n",
    "terapi_df[\"insert_date\"] =  pd.to_datetime(terapi_df['insert'], format='%Y-%m-%d',errors = \"coerce\")\n",
    "terapi_df[\"end_date\"] =  pd.to_datetime(terapi_df['exposed'], format='%Y-%m-%d',errors = \"coerce\")\n",
    "\n",
    "demographic_addinfo_df[\"terapi drug\"] = demographic_addinfo_df.parallel_apply(lambda x: get_terapi_drug(x[\"patient code\"],x[\"visit_date\"]),axis=1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e7f9dd70",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Adding drug classification to the data\n",
    "drug_classification_df = pd.read_excel(\"../drug_classification.xlsx\")\n",
    "drug_types = [column_name for column_name in drug_classification_df.columns.to_list() if column_name != \"drugs\"]\n",
    "drug_type_catagory = {drug_type:drug_types.index(drug_type) for drug_type in drug_types}\n",
    "no_treatment_label = len(drug_type_catagory)\n",
    "drug_type_catagory[\"no_treatment\"] = no_treatment_label\n",
    "drug_type_catagory_inv = {drug_type_catagory[entry]:entry for entry in drug_type_catagory}\n",
    "\n",
    "\n",
    "drug_catagorized = {}\n",
    "for index,row in drug_classification_df.iterrows():\n",
    "    row_index_of_drug_type = row[row == 1].index\n",
    "    if len(row_index_of_drug_type.to_list()) > 1:\n",
    "        drug_catagorized[row[\"drugs\"]] = [drug_type_catagory[entry] for entry in row_index_of_drug_type]\n",
    "    else:\n",
    "        drug_catagorized[row[\"drugs\"]] = drug_type_catagory[row_index_of_drug_type[0]]\n",
    "drug_catagorized['Ingen behandling'] = no_treatment_label\n",
    "drug_catagorized['no treatment'] = no_treatment_label"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "fdd9d827",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Adding drug infor the df\n",
    "for index,row in tqdm.tqdm(demographic_addinfo_df.iterrows()):\n",
    "    entry = row[\"terapi drug\"]\n",
    "    if type(entry) == list:\n",
    "        for drugname in entry:\n",
    "            if drugname in drug_catagorized:\n",
    "                if type(drug_catagorized[drugname]) == list:\n",
    "                    for drug_label in drug_catagorized[drugname]:\n",
    "                        rowname = drug_type_catagory_inv[drug_label]\n",
    "                        demographic_addinfo_df.loc[index,rowname]= 1\n",
    "                else:\n",
    "                    rowname = drug_type_catagory_inv[drug_catagorized[drugname]]\n",
    "                    #drug_catagory = drug_catagorized[drugname]\n",
    "                    demographic_addinfo_df.loc[index,rowname]= 1\n",
    "            else:\n",
    "                demographic_addinfo_df.loc[index,\"not found\"]= drugname\n",
    "    else:\n",
    "        if entry in drug_catagorized:\n",
    "            if type(drug_catagorized[entry]) == list:\n",
    "                for drug_label in drug_catagorized[entry]:\n",
    "                    rowname = drug_type_catagory_inv[drug_label]\n",
    "                    demographic_addinfo_df.loc[index,rowname]= 1\n",
    "            else:\n",
    "                rowname = drug_type_catagory_inv[drug_catagorized[entry]]\n",
    "                #drug_catagory = drug_catagorized[entry]\n",
    "                demographic_addinfo_df.loc[index,rowname]= 1\n",
    "        else:\n",
    "            demographic_addinfo_df.loc[index,\"not found\"]= entry\n",
    "            \n",
    "# Fill 0 in case of NA\n",
    "demographic_addinfo_df[[\"first_line_DMT\",\"second_line_DMT\",\"no_treatment\",\"relapse_treatment_drugs\",\"other_drugs\",\"stem_cell_treatment\"]] = demographic_addinfo_df[[\"first_line_DMT\",\"second_line_DMT\",\"no_treatment\",\"relapse_treatment_drugs\",\"other_drugs\",\"stem_cell_treatment\"]].fillna(0)\n",
    "# if no treatment and some treatment is present. Set no treatment to 0\n",
    "index_to_have_no_treatment_changed = (demographic_addinfo_df[\"no_treatment\"] == 1) & ((demographic_addinfo_df[\"first_line_DMT\"] == 1)|(demographic_addinfo_df[\"second_line_DMT\"] == 1)|(demographic_addinfo_df[\"relapse_treatment_drugs\"] == 1)|(demographic_addinfo_df[\"other_drugs\"] == 1))\n",
    "demographic_addinfo_df.loc[index_to_have_no_treatment_changed,\"no_treatment\"] = 0"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "fa257f22",
   "metadata": {},
   "source": [
    "# Relapse"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b250df89",
   "metadata": {},
   "outputs": [],
   "source": [
    "skov_df = pd.read_csv(filepath + \"skov.csv\")\n",
    "df = copy.deepcopy(demographic_addinfo_df)\n",
    "skov_df[\"date\"] =  pd.to_datetime(skov_df['date'], format='%Y-%m-%d')\n",
    "skov_df = skov_df.dropna(subset=[\"date\"]) # Remove if date is not present\n",
    "skov_df = skov_df.sort_values(\"date\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "31ccf2a0",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Remove duplicate debut date\n",
    "# We can see that there are multiple debut date for the same patient\n",
    "# We need to change the second entry of the debut to first one.\n",
    "debut_pcode = skov_df[skov_df[\"debut relapse\"] == 1][\"patient code\"].to_list() # Find pcode with debut == 1\n",
    "duplicated_pcodes = Counter(debut_pcode)\n",
    "duplicated_pcodes = [entry for entry in duplicated_pcodes if duplicated_pcodes[entry] > 1] # Find the duplicate pcodes\n",
    "\n",
    "\n",
    "# Change the debut relapse to 1 for the first date, when there are duplicates for the same patient\n",
    "for pcode in duplicated_pcodes:\n",
    "    subset_df = skov_df[(skov_df[\"patient code\"] == pcode) & (skov_df[\"debut relapse\"] == 1)]\n",
    "    first_debut_date_index = subset_df.sort_values(\"date\")[\"date\"].index[0]\n",
    "    skov_df.loc[subset_df.index,\"debut relapse\"] = 0\n",
    "    skov_df.loc[first_debut_date_index,\"debut relapse\"] = 1\n",
    "    \n",
    "# Remove patient entry if there is a relapse before debut date\n",
    "for pcode in list(set(skov_df[\"patient code\"].to_list())):\n",
    "    sub_df = skov_df[skov_df[\"patient code\"] == pcode]\n",
    "    if 1 in sub_df[\"debut relapse\"].to_list():\n",
    "        index_of_1 = sub_df[\"debut relapse\"].to_list().index(1)\n",
    "        if index_of_1 != 0:\n",
    "            index_to_remove = sub_df[\"debut relapse\"].index[index_of_1]\n",
    "            skov_df.drop(index=index_to_remove)\n",
    "            \n",
    "print (\"Number of patients with debut date =\",len(set(skov_df[skov_df[\"debut relapse\"] == 1][\"patient code\"].to_list())))\n",
    "print (\"Total number of patients\",len(set(skov_df[\"patient code\"].to_list())))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9b20f9d1",
   "metadata": {},
   "outputs": [],
   "source": [
    "# If debut relapse is not found. Consider the first date as the debut relapse\n",
    "for pcode in list(set(skov_df[\"patient code\"].to_list())):\n",
    "    sub_df = skov_df[skov_df[\"patient code\"] == pcode]\n",
    "    if 1 not in sub_df[\"debut relapse\"].to_list():\n",
    "        skov_df.loc[sub_df.index[0],\"debut relapse\"] = 1\n",
    "print (\"Number of patients with debut date =\",len(set(skov_df[skov_df[\"debut relapse\"] == 1][\"patient code\"].to_list())))\n",
    "print (\"Total number of patients\",len(set(skov_df[\"patient code\"].to_list())))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7b760ee4",
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_multi_focal_skov(entry):\n",
    "    if entry == 0.0:\n",
    "        return 1\n",
    "    else:\n",
    "        return np.nan\n",
    "\n",
    "skov_df[\"multi_focal\"] = skov_df[\"mono on\"].apply(lambda x: get_multi_focal_skov(x))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3a5993a9",
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "sub_df = skov_df.fillna(0) # Fill all the na values with 0\n",
    "pandarallel.initialize(progress_bar=True,nb_workers=32)\n",
    "def get_number_of_relapse_using_visitdate(visit_date,patient_code):\n",
    "    df = sub_df[sub_df[\"patient code\"]==patient_code][sub_df[sub_df[\"patient code\"]==patient_code][\"date\"] <= visit_date]\n",
    "    if len(df) > 0:\n",
    "        mono_on_sum = sum(df[\"mono on\"])\n",
    "        multi_focal_sum = sum(df[\"multi_focal\"])\n",
    "        monofocal_sum = sum(df[\"monofocal\"])\n",
    "        afferent_non_on_sum = sum(df[\"afferent non on\"])\n",
    "        steroid_treatment_sum = sum(df[\"steroid treatment\"])\n",
    "        is_last_relapse_steroid_treated = df[\"steroid treatment\"].to_list()[-1]\n",
    "        is_last_relapse_completely_remitted = df[\"complete remission\"].to_list()[-1]\n",
    "        last_relapse_date = df[\"date\"].to_list()[-1]\n",
    "        debut_relapse_date = df[\"date\"].to_list()[0]\n",
    "        \n",
    "        return pd.Series([mono_on_sum,multi_focal_sum, monofocal_sum, afferent_non_on_sum, steroid_treatment_sum, is_last_relapse_steroid_treated, is_last_relapse_completely_remitted, last_relapse_date, debut_relapse_date])\n",
    "    return pd.Series([np.nan,np.nan,np.nan,np.nan,np.nan,np.nan,np.nan,np.nan])\n",
    "\n",
    "demographic_addinfo_df[[\"mono_on_sum\",\"multi_focal_sum\",\"monofocal_sum\",\"afferent_non_on_sum\",\"steroid_treatment_sum\",\"is_last_relapse_steroid_treated\",\"is_last_relapse_completely_remitted\",\"last_relapse_date\",\"debut_relapse_date\"]] = demographic_addinfo_df.parallel_apply(lambda x: get_number_of_relapse_using_visitdate(x[\"visit_date\"],x[\"patient code\"]),axis=1)\n",
    "demographic_addinfo_df[\"age_at_relapse\"] = np.round((demographic_addinfo_df[\"last_relapse_date\"]-demographic_addinfo_df[\"birthday\"]) / np.timedelta64(1, 'Y'),3)\n",
    "demographic_addinfo_df[\"age_at_debut_relapse\"] = np.round((demographic_addinfo_df[\"debut_relapse_date\"]-demographic_addinfo_df[\"birthday\"]) / np.timedelta64(1, 'Y'),3)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5aa1c629",
   "metadata": {},
   "source": [
    "# MRI"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "bc338406",
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "mri_df = pd.read_csv(filepath + \"mri.csv\")\n",
    "mri_df[\"date\"] =  pd.to_datetime(mri_df['date'], format='%Y-%m-%d', errors = 'coerce')\n",
    "mri_df = mri_df[mri_df[\"date\"].notna()]\n",
    "mri_df= mri_df.sort_values(by=[\"date\"])"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f553ccc7",
   "metadata": {},
   "source": [
    "## Column definition\n",
    "Date\n",
    "\n",
    "\"total number of t2 lesions\"  range of ['3', '1-9', '>20', '6', '5', '1', '2', '7', '10-20', '0', '8', '4', '9']  \n",
    "\"number of new or advanced t2 lesions\"  range of ['3', '1-9', '>20', '6', '5', '1', '2', '7', '10-20', '0', '8', '4', '9']  \n",
    "\"total number of lesions with brain barrier injury\" range of ['3', '>20', '6', '5', '1', '2', '7', '10-20', '0', '8', '4', '9']  \n",
    "\"total number of lesions with spinal cord barrier injury\" range of ['3', 0.0, 2.0, 1.0, 3.0, 5.0, 6.0, '>20', '6', '5', '1', '2', '7', '10-20', '0', '8', '4', '9']  \n",
    "\"number of lesions compared to previous spinal cord\" range of  [1.0, 2.0, 3.0]  \n",
    "  \n",
    "\"total brain volume\" in millilitre  \n",
    "\"bpf\" is in percentage  \n",
    "  \n",
    "\"lesion in the spinal cord\" is binary [0,1]  "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1fec4319-a1ce-48f0-a9da-8752581fca23",
   "metadata": {},
   "outputs": [],
   "source": [
    "def lesion_count_to_catagory(input_value):\n",
    "    if input_value >= 20:\n",
    "        input_value = 20\n",
    "    if input_value < 20 and input_value >= 10:\n",
    "        input_value = 20\n",
    "    \n",
    "    return lesion_catagory[lesion_count_dict_inv[input_value]]\n",
    "\n",
    "def mri_lesion_count_to_int_number(input_value):\n",
    "    try: # if int value\n",
    "        return int(input_value)\n",
    "    except:\n",
    "        try: # 1-9,10-20 or >20\n",
    "            return lesion_count_dict[input_value]\n",
    "        except: # If nan\n",
    "            return 0"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "467e5cdc",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Making catagories for lesion count\n",
    "lesion_catagory = {}\n",
    "for entry in [0.0, 1.0, 2.0, 3.0, 5.0, 6.0, '6', '9', '3', '5', '2', '0', '1-9', '>20', '1', '7', '4', '10-20', '8']:\n",
    "    try:\n",
    "        if int(entry) == 0:\n",
    "            lesion_catagory[int(entry)] = 0\n",
    "        else:\n",
    "            lesion_catagory[int(entry)] = 1\n",
    "    except:\n",
    "        if entry == '1-9':\n",
    "            lesion_catagory[entry] = 1\n",
    "        if entry == '10-20':\n",
    "            lesion_catagory[entry] = 2\n",
    "        if entry == '>20':\n",
    "            lesion_catagory[entry] = 3\n",
    "            \n",
    "lesion_count_dict = {}\n",
    "for entry in lesion_catagory:\n",
    "    try:\n",
    "        entry = int(entry)\n",
    "        lesion_count_dict[entry] = entry\n",
    "    except:\n",
    "        if entry == '1-9':\n",
    "            lesion_count_dict[entry] = 9\n",
    "        if entry == '10-20':\n",
    "            lesion_count_dict[entry] = 20\n",
    "        if entry == '>20':\n",
    "            lesion_count_dict[entry] = 25\n",
    "            \n",
    "lesion_count_dict_inv = {lesion_count_dict[entry]:entry for entry in lesion_count_dict}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "42d3295a",
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "def process_mri_data(visit_date,patient_code):\n",
    "    df = mri_df[mri_df[\"patient code\"]==patient_code][mri_df[mri_df[\"patient code\"]==patient_code][\"date\"] <= visit_date]\n",
    "    df= df.sort_values(by=[\"date\"])\n",
    "    if len(df) > 0:\n",
    "        t2_lesion_count = mri_lesion_count_to_int_number(df[\"total number of t2 lesions\"].to_list()[-1])\n",
    "        t2_new_lesion_count = mri_lesion_count_to_int_number(df[\"number of new or advanced t2 lesions\"].to_list()[-1])\n",
    "        t2_lesion_catagory = lesion_count_to_catagory(t2_lesion_count+t2_new_lesion_count)\n",
    "        \n",
    "        brain_barrier_lesion_count = mri_lesion_count_to_int_number(df[\"total number of lesions with brain barrier injury\"].to_list()[-1])\n",
    "        brain_barrier_lesion_catagory = lesion_count_to_catagory(brain_barrier_lesion_count)\n",
    "        \n",
    "        spinal_barrier_lesion_count = mri_lesion_count_to_int_number(df[\"total number of lesions with spinal cord barrier injury\"].to_list()[-1])\n",
    "        spinal_barrier_lesion_catagory = lesion_count_to_catagory(spinal_barrier_lesion_count)\n",
    "        \n",
    "        checkup_date = df[\"date\"].to_list()[-1]\n",
    "        \n",
    "        \n",
    "        return pd.Series([t2_lesion_catagory,brain_barrier_lesion_catagory,spinal_barrier_lesion_catagory,checkup_date])\n",
    "    return pd.Series([np.nan,np.nan,np.nan,np.nan])\n",
    "\n",
    "demographic_addinfo_df[[\"t2_lesion_catagory\",\"brain_barrier_lesion_catagory\",\"spinal_barrier_lesion_catagory\",\"mri_date\"]] = demographic_addinfo_df.parallel_apply(lambda x: process_mri_data(x[\"visit_date\"],x[\"patient code\"]),axis=1)\n",
    "demographic_addinfo_df[\"age_at_mri\"] = np.round((demographic_addinfo_df[\"mri_date\"]-demographic_addinfo_df[\"birthday\"]) / np.timedelta64(1, 'Y'),3)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "66c2c8a7",
   "metadata": {},
   "source": [
    "# MSIS-29"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f198d119",
   "metadata": {},
   "outputs": [],
   "source": [
    "selected_columns = ['date','msis_01', 'msis_02', 'msis_03', 'msis_04', 'msis_05', 'msis_06', 'msis_07', \n",
    " 'msis_08', 'msis_09', 'msis_10', 'msis_11', 'msis_12', 'msis_13', 'msis_14', 'msis_15', \n",
    " 'msis_16', 'msis_17', 'msis_18', 'msis_19', 'msis_20', 'msis_21', 'msis_22', 'msis_23', \n",
    " 'msis_24', 'msis_25', 'msis_26', 'msis_27', 'msis_28', 'msis_29', 'physically', 'psychologically', 'physically_100', 'psychologically_100']\n",
    "\n",
    "selected_columns_renamed = ['msis_date','msis_01', 'msis_02', 'msis_03', 'msis_04', 'msis_05', 'msis_06', 'msis_07', \n",
    " 'msis_08', 'msis_09', 'msis_10', 'msis_11', 'msis_12', 'msis_13', 'msis_14', 'msis_15', \n",
    " 'msis_16', 'msis_17', 'msis_18', 'msis_19', 'msis_20', 'msis_21', 'msis_22', 'msis_23', \n",
    " 'msis_24', 'msis_25', 'msis_26', 'msis_27', 'msis_28', 'msis_29', 'msis_physically', 'msis_psychologically', 'msis_physically_100', 'msis_psychologically_100']"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "37dce195",
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "msis_df = pd.read_csv(filepath + \"msis_29.csv\")\n",
    "msis_df.columns = [entry.replace(\" \",\"_\") for entry in msis_df.columns.to_list()] # replace column names\n",
    "msis_df[\"date\"] =  pd.to_datetime(msis_df['date'], format='%Y-%m-%d',errors = \"coerce\")\n",
    "msis_df = msis_df[msis_df[\"date\"].notna()]\n",
    "msis_df= msis_df.sort_values(by=[\"date\"])\n",
    "\n",
    "\n",
    "def get_last_msis_using_date(visit_date,patient_code):\n",
    "    df = msis_df[msis_df[\"patient_code\"]==patient_code][msis_df[msis_df[\"patient_code\"]==patient_code][\"date\"] <= visit_date]\n",
    "    df= df.sort_values(\"date\")\n",
    "    if len(df) > 0:\n",
    "        last_msis_date = df[\"date\"].to_list()[-1]\n",
    "        output_values = df[df[\"date\"] == last_msis_date][selected_columns].values.tolist()[0]\n",
    "        return pd.Series(output_values)\n",
    "    return pd.Series([np.nan for entry in selected_columns])\n",
    "\n",
    "demographic_addinfo_df[selected_columns_renamed] = demographic_addinfo_df.parallel_apply(lambda x: get_last_msis_using_date(x[\"visit_date\"],x[\"patient code\"]),axis=1)\n",
    "demographic_addinfo_df[\"age_at_msis\"] = np.round((demographic_addinfo_df[\"msis_date\"]-demographic_addinfo_df[\"birthday\"]) / np.timedelta64(1, 'Y'),3)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e031a4dc",
   "metadata": {},
   "source": [
    "# Cleaning and saving the whole data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0a1f79d9",
   "metadata": {},
   "outputs": [],
   "source": [
    "df = demographic_addinfo_df.copy(deep=True)\n",
    "df= df.sort_values(by=[\"patient code\",'visit_date'])\n",
    "# Snapshot data\n",
    "\n",
    "df = df.reset_index()\n",
    "df = df.drop([\"index\"],axis=1)\n",
    "df[\"sex_label\"] = (df[\"sex\"] == \"MAN\").astype(\"int\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a8842c65-f142-41c4-8994-54ce1cafce38",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Remove entries when debut date is after visit date\n",
    "print (\"Entries removed =\", np.sum((((df[\"visit_date\"] - df[\"debut_date\"])/ np.timedelta64(1, 'D')) < 0)))\n",
    "df = df[~(((df[\"visit_date\"] - df[\"debut_date\"])/ np.timedelta64(1, 'D')) < 0)]\n",
    "\n",
    "# Remove entries if birthday is bigger than visit date\n",
    "print (\"Entries removed = \",np.sum(((df[\"birthday\"] - df[\"visit_date\"])/ np.timedelta64(1, 'D')) >= 0))\n",
    "df = df[~(((df[\"birthday\"] - df[\"visit_date\"])/ np.timedelta64(1, 'D')) >= 0)]\n",
    "\n",
    "# Remove entries where last relapse date is before debut relapse date\n",
    "print (\"Entries removed = \",np.sum(((df[\"debut_relapse_date\"] - df[\"last_relapse_date\"])/ np.timedelta64(1, 'D')) > 0))\n",
    "df = df[~(((df[\"debut_relapse_date\"] - df[\"last_relapse_date\"])/ np.timedelta64(1, 'D')) > 0)]\n",
    "\n",
    "# Fixing debut date\n",
    "df[\"revised_debut_date\"] = df[\"diagnosis_date\"]\n",
    "df.loc[df[\"diagnosis_date\"].isna(),\"revised_debut_date\"] = df[\"debut_date\"]\n",
    "df[\"revised_debut_age\"] = np.round((df[\"revised_debut_date\"]-df[\"birthday\"]) / np.timedelta64(1, 'Y'),3)\n",
    "\n",
    "# If the values for these features are not available, it will be imputed with padding index (-1)\n",
    "other_columns = ['no_treatment', 'first_line_DMT', 'second_line_DMT', 'other_drugs', 'relapse_treatment_drugs','stem_cell_treatment',\n",
    "                     'eq5d_score','age_at_eq5d','sdmt_score','age_at_sdmt',\n",
    "                     'mono_on_sum','monofocal_sum','multi_focal_sum','afferent_non_on_sum','steroid_treatment_sum','is_last_relapse_steroid_treated','is_last_relapse_completely_remitted','age_at_relapse','age_at_debut_relapse',\n",
    "                't2_lesion_catagory', 'brain_barrier_lesion_catagory', 'spinal_barrier_lesion_catagory','age_at_mri',\n",
    "                     'msis_01', 'msis_02', 'msis_03', 'msis_04', 'msis_05', 'msis_06', 'msis_07', 'msis_08', 'msis_09', \n",
    "                'msis_10', 'msis_11', 'msis_12', 'msis_13', 'msis_14', 'msis_15', 'msis_16', 'msis_17', 'msis_18', 'msis_19',\n",
    "                'msis_20', 'msis_21', 'msis_22', 'msis_23', 'msis_24', 'msis_25', 'msis_26', 'msis_27', 'msis_28', 'msis_29', \n",
    "                'msis_physically', 'msis_psychologically', 'msis_physically_100', 'msis_psychologically_100', 'age_at_msis','sex_label'] \n",
    "\n",
    "base_columns = ['edss value/score','age_at_visit','revised_debut_age'] # If values for there features are not present, the data is removed\n",
    "\n",
    "df[other_columns] = df[other_columns].fillna(-1)\n",
    "\n",
    "df.to_csv(\"data_for_pca_python_cleaned.csv\",index=False)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8332b29d-76dc-48c9-80b5-b5239159e715",
   "metadata": {},
   "source": [
    "# Creating data splits (Train/Valid/Test/Calibration)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "aa71af6c-0fcc-4449-9ebc-c86de80eb466",
   "metadata": {},
   "outputs": [],
   "source": [
    "# select only RRMS and SPMS patients\n",
    "# remove ppms patients\n",
    "df = df[df[\"progress during visit\"] != \"PP\"]\n",
    "df.loc[df[\"progress during visit\"] == \"RR\",\"y_label\"]= 0\n",
    "df.loc[df[\"progress during visit\"] == \"SP\",\"y_label\"]= 1\n",
    "\n",
    "# Split to train, valid and test data\n",
    "# The split data is kept as object to enable augmentation in the training data\n",
    "train_size=0.65\n",
    "calibration_size = 0.05\n",
    "valid_size=0.05\n",
    "test_size=1 - (train_size+valid_size+calibration_size)\n",
    "random_state=0"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d7dc23d3-7a13-4841-9867-561293e47094",
   "metadata": {},
   "outputs": [],
   "source": [
    "# segregate patients into different catagories for stratification during data splitting\n",
    "def get_RRMS_SPMS_transitioning_patients(input_df):\n",
    "    pcode_label_dict = {}\n",
    "    for entry in list(df[[\"patient code\",\"y_label\"]].values):\n",
    "        pcode = int(entry[0])\n",
    "        label = int(entry[1])\n",
    "        if pcode in pcode_label_dict:\n",
    "            pcode_label_dict[pcode].append(label)\n",
    "        else:\n",
    "            pcode_label_dict[pcode] = [label]\n",
    "    pcode_label_dict = {entry:list(set(pcode_label_dict[entry])) for entry in pcode_label_dict}\n",
    "    \n",
    "    rr_pcode,sp_pcode,t_pcode = [],[],[]\n",
    "    for pcode in pcode_label_dict:\n",
    "        if len(pcode_label_dict[pcode]) == 2:\n",
    "            t_pcode.append([pcode,2])\n",
    "        elif pcode_label_dict[pcode] == [0]:\n",
    "            rr_pcode.append([pcode,0])\n",
    "        elif pcode_label_dict[pcode] == [1]:\n",
    "            sp_pcode.append([pcode,1])\n",
    "\n",
    "\n",
    "    rr_pcode.extend(sp_pcode)\n",
    "    rr_pcode.extend(t_pcode)\n",
    "    return rr_pcode\n",
    "\n",
    "def create_data_splits(input_df,train_size=0.7,valid_size=0.15,test_size=0.15,calibration_set_size=0.1,random_state=0):\n",
    "    x_train, x_valid_test_calibration = train_test_split(all_pcodes, test_size=(valid_size+test_size+calibration_size),random_state=random_state,stratify =all_pcodes[:,1:])\n",
    "    x_calibration, x_valid_test = train_test_split(x_valid_test_calibration, test_size=(valid_size+test_size)/(1-train_size),random_state=random_state,stratify=x_valid_test_calibration[:,1:])\n",
    "    x_valid, x_test = train_test_split(x_valid_test, test_size=test_size/(1-(train_size+calibration_size)) ,random_state=random_state,stratify=x_valid_test[:,1:])\n",
    "    \n",
    "    train_pcodes = x_train[:,:1].squeeze()\n",
    "    valid_pcodes = x_valid[:,:1].squeeze()\n",
    "    test_pcodes = x_test[:,:1].squeeze()\n",
    "    calibration_pcodes = x_calibration[:,:1].squeeze()\\\n",
    "\n",
    "    train_df = df[df[\"patient code\"].isin(train_pcodes)]\n",
    "    valid_df = df[df[\"patient code\"].isin(valid_pcodes)]\n",
    "    test_df = df[df[\"patient code\"].isin(test_pcodes)]\n",
    "    calibration_df = df[df[\"patient code\"].isin(calibration_pcodes)]\n",
    "\n",
    "\n",
    "    return train_df,valid_df,test_df,calibration_df"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7e22512f-f9b7-4986-8a63-dbbc8a5ab3ef",
   "metadata": {},
   "outputs": [],
   "source": [
    "all_pcodes = np.array(get_RRMS_SPMS_transitioning_patients(df))\n",
    "train_df,valid_df,test_df,calibration_df= create_data_splits(df,train_size=train_size,valid_size=valid_size, test_size=test_size,calibration_set_size=calibration_size,random_state=random_state)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c5c361cc-8e1e-4367-b69a-0fbc67e0a0e6",
   "metadata": {},
   "outputs": [],
   "source": [
    "os.mkdir(\"data_splits\")\n",
    "train_df.to_csv(\"data_splits/train.csv\",index=False)\n",
    "valid_df.to_csv(\"data_splits/valid.csv\",index=False)\n",
    "test_df.to_csv(\"data_splits/test.csv\",index=False)\n",
    "calibration_df.to_csv(\"data_splits/calibration.csv\",index=False)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b58af520",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "31a0080b",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4c0344d8",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f2d4837d",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f530e838-b3ec-4ef7-84e3-d5fa15ebc3fa",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python [conda env:base] *",
   "language": "python",
   "name": "conda-base-py"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.10.12"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}