pc_demo_02.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "e32eca0a-aed8-47d7-9816-4d4c57bf5871",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Objectives:\n",
    "# 1. Include gene expression features in training after preprocessing (binning, normilization).\n",
    "# 2. Save encoders and other preprocessing objects to be used in validation examples.\n",
    "# 3. Feature selection"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "5a48b651-8003-4462-b12c-4e394404fecd",
   "metadata": {},
   "outputs": [],
   "source": [
    "# conda install pytorch torchvision torchaudio cpuonly -c pytorch\n",
    "# conda install conda-forge::polars\n",
    "# conda install conda-forge::xgboost"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "f361a590-36cc-4bdb-b7a5-253a460efa06",
   "metadata": {},
   "outputs": [],
   "source": [
    "import polars as pl\n",
    "import numpy as np\n",
    "import torch\n",
    "import torch.nn as nn\n",
    "import torch.optim as optim\n",
    "import xgboost as xgb\n",
    "from torch.utils.data import DataLoader, Dataset\n",
    "from sklearn.model_selection import train_test_split, KFold\n",
    "from sklearn.preprocessing import LabelEncoder, StandardScaler, MinMaxScaler\n",
    "from sklearn.model_selection import GridSearchCV, GroupKFold\n",
    "from sklearn.metrics import label_ranking_loss\n",
    "from sklearn.datasets import make_regression\n",
    "from tqdm import tqdm\n",
    "from mrmr import mrmr_regression\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "001f2c5f-a72c-421c-800a-1d406d93bef6",
   "metadata": {},
   "outputs": [],
   "source": [
    "target = 'IC50_PUBLISHED'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "278a1f0e-f5f6-4c8b-ab50-4340f4766b76",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Shape of df_dose_resp_gdsc2 = (118908, 3)\n",
      "Unique cell lines (ARXSPAN_ID) = 794\n",
      "Unique drugs = 175\n",
      "Unique combinations of cell line x drug = 116377\n",
      "<bound method DataFrame.head of shape: (118_908, 3)\n",
      "┌─────────┬────────────┬────────────────┐\n",
      "│ DRUG_ID ┆ ARXSPAN_ID ┆ IC50_PUBLISHED │\n",
      "│ ---     ┆ ---        ┆ ---            │\n",
      "│ i64     ┆ str        ┆ f64            │\n",
      "╞═════════╪════════════╪════════════════╡\n",
      "│ 1003    ┆ ACH-000958 ┆ 0.025129       │\n",
      "│ 1003    ┆ ACH-000651 ┆ 0.049577       │\n",
      "│ 1003    ┆ ACH-000856 ┆ 0.028549       │\n",
      "│ 1003    ┆ ACH-000360 ┆ 0.039996       │\n",
      "│ 1003    ┆ ACH-001199 ┆ 1.986678       │\n",
      "│ …       ┆ …          ┆ …              │\n",
      "│ 2172    ┆ ACH-000288 ┆ 25.410793      │\n",
      "│ 2172    ┆ ACH-001065 ┆ 0.339325       │\n",
      "│ 2172    ┆ ACH-000930 ┆ 7.780877       │\n",
      "│ 2172    ┆ ACH-000859 ┆ 534.688321     │\n",
      "│ 2172    ┆ ACH-000536 ┆ 120.177282     │\n",
      "└─────────┴────────────┴────────────────┘>\n"
     ]
    }
   ],
   "source": [
    "# Read cell line to drug id to ic50\n",
    "# https://depmap.org/portal/data_page/?tab=allData&releasename=Sanger+GDSC1+and+GDSC2&filename=sanger-dose-response.csv\n",
    "# IC50 score of drugs (DRUG_ID) per cell line (COSMIC_ID) for GDSC1 and GDSC2 \n",
    "try:\n",
    "    df_dose_resp = pl.read_csv(\"C:\\\\Users\\\\chris\\\\rank-rx\\\\data\\\\sanger-dose-response.csv\")\n",
    "    df_dose_resp_gdsc2 = df_dose_resp.filter(pl.col(\"DATASET\") == \"GDSC2\")\n",
    "    df_dose_resp_gdsc2_edited = df_dose_resp_gdsc2.select([\"DRUG_ID\", \"ARXSPAN_ID\", \"IC50_PUBLISHED\"])\n",
    "    print(\"Shape of df_dose_resp_gdsc2 = {}\".format(df_dose_resp_gdsc2_edited.shape))\n",
    "    print(\"Unique cell lines (ARXSPAN_ID) = {}\".format(df_dose_resp_gdsc2_edited['ARXSPAN_ID'].unique().len()))\n",
    "    print(\"Unique drugs = {}\".format(df_dose_resp_gdsc2_edited['DRUG_ID'].unique().len()))\n",
    "    grouped = df_dose_resp_gdsc2_edited.group_by(['ARXSPAN_ID', 'DRUG_ID']).agg(pl.len())\n",
    "    print(\"Unique combinations of cell line x drug = {}\".format(grouped.shape[0]))\n",
    "    print(df_dose_resp_gdsc2_edited.head)\n",
    "except Exception as e:\n",
    "    print(f\"Error: {e}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "ddeec08c-c054-4214-a865-35eb5a97d69b",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Shape of df_depmap_model = (1959, 6)\n",
      "Unique cell lines (ModelID) ACH-XXXXXX = 1959\n",
      "shape: (5, 6)\n",
      "┌────────────┬──────────────┬─────────────┬────────┬─────────────┬─────────────────────┐\n",
      "│ ModelID    ┆ OncotreeCode ┆ AgeCategory ┆ Sex    ┆ PatientRace ┆ PrimaryOrMetastasis │\n",
      "│ ---        ┆ ---          ┆ ---         ┆ ---    ┆ ---         ┆ ---                 │\n",
      "│ str        ┆ str          ┆ str         ┆ str    ┆ str         ┆ str                 │\n",
      "╞════════════╪══════════════╪═════════════╪════════╪═════════════╪═════════════════════╡\n",
      "│ ACH-000001 ┆ HGSOC        ┆ Adult       ┆ Female ┆ caucasian   ┆ Metastatic          │\n",
      "│ ACH-000002 ┆ AML          ┆ Adult       ┆ Female ┆ caucasian   ┆ Primary             │\n",
      "│ ACH-000003 ┆ COAD         ┆ Adult       ┆ Male   ┆ caucasian   ┆ Primary             │\n",
      "│ ACH-000004 ┆ AML          ┆ Adult       ┆ Male   ┆ caucasian   ┆ Primary             │\n",
      "│ ACH-000005 ┆ AML          ┆ Adult       ┆ Male   ┆ caucasian   ┆ null                │\n",
      "└────────────┴──────────────┴─────────────┴────────┴─────────────┴─────────────────────┘\n"
     ]
    }
   ],
   "source": [
    "# Read demographics and cancer type\n",
    "# https://depmap.org/portal/data_page/?tab=allData&releasename=DepMap+Public+24Q2&filename=Model.csv\n",
    "# Mapping between 'ModelID', 'PatientID', 'SangerModelID', 'COSMICID', etc\n",
    "# ModelID here is the cell line id.\n",
    "# OncotreeCode is the type of cancer.\n",
    "try:\n",
    "    df_depmap_model = pl.read_csv(\"C:\\\\Users\\\\chris\\\\rank-rx\\\\data\\\\Model.csv\")\n",
    "    df_depmap_model_edited = df_depmap_model.select(['ModelID', 'OncotreeCode', 'AgeCategory', 'Sex', 'PatientRace', 'PrimaryOrMetastasis'])\n",
    "    print(\"Shape of df_depmap_model = {}\".format(df_depmap_model_edited.shape))\n",
    "    print(\"Unique cell lines (ModelID) ACH-XXXXXX = {}\".format(df_depmap_model_edited['ModelID'].unique().len()))\n",
    "    print(df_depmap_model_edited.head())\n",
    "except Exception as e:\n",
    "    print(f\"Error: {e}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "84ba5e84-2171-43e5-b385-233d27b5bcae",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Shape of df_gene_express = (1517, 19138)\n",
      "shape: (5, 19_138)\n",
      "┌───────────┬───────────┬───────────┬───────────┬───┬───────────┬───────────┬───────────┬──────────┐\n",
      "│ ARXSPAN_I ┆ ZNF891    ┆ ARMC10    ┆ PTGER4    ┆ … ┆ DNMT3B    ┆ ZCCHC10   ┆ PRSS2     ┆ ADAMTSL4 │\n",
      "│ D         ┆ (10106020 ┆ (83787)   ┆ (5734)    ┆   ┆ (1789)    ┆ (54819)   ┆ (5645)    ┆ (54507)  │\n",
      "│ ---       ┆ 0)        ┆ ---       ┆ ---       ┆   ┆ ---       ┆ ---       ┆ ---       ┆ ---      │\n",
      "│ str       ┆ ---       ┆ f64       ┆ f64       ┆   ┆ f64       ┆ f64       ┆ f64       ┆ f64      │\n",
      "│           ┆ f64       ┆           ┆           ┆   ┆           ┆           ┆           ┆          │\n",
      "╞═══════════╪═══════════╪═══════════╪═══════════╪═══╪═══════════╪═══════════╪═══════════╪══════════╡\n",
      "│ ACH-00005 ┆ 0.984137  ┆ 4.524944  ┆ 2.019524  ┆ … ┆ 2.320999  ┆ 5.005448  ┆ 0.169594  ┆ 1.356288 │\n",
      "│ 8         ┆           ┆           ┆           ┆   ┆           ┆           ┆           ┆          │\n",
      "│ ACH-00008 ┆ 1.335101  ┆ 3.974294  ┆ 1.923885  ┆ … ┆ 2.48262   ┆ 4.645063  ┆ 0.229243  ┆ 0.795239 │\n",
      "│ 2         ┆           ┆           ┆           ┆   ┆           ┆           ┆           ┆          │\n",
      "│ ACH-00086 ┆ 0.752579  ┆ 4.698468  ┆ 4.105001  ┆ … ┆ 2.047717  ┆ 4.499003  ┆ 0.26387   ┆ 6.488749 │\n",
      "│ 7         ┆           ┆           ┆           ┆   ┆           ┆           ┆           ┆          │\n",
      "│ ACH-00043 ┆ 1.589078  ┆ 5.694567  ┆ 3.362272  ┆ … ┆ 2.239845  ┆ 4.778812  ┆ 0.26387   ┆ 1.205591 │\n",
      "│ 7         ┆           ┆           ┆           ┆   ┆           ┆           ┆           ┆          │\n",
      "│ ACH-00054 ┆ 0.530802  ┆ 5.76315   ┆ 0.609288  ┆ … ┆ 1.300537  ┆ 4.56258   ┆ -0.012068 ┆ 0.707417 │\n",
      "│ 4         ┆           ┆           ┆           ┆   ┆           ┆           ┆           ┆          │\n",
      "└───────────┴───────────┴───────────┴───────────┴───┴───────────┴───────────┴───────────┴──────────┘\n"
     ]
    }
   ],
   "source": [
    "# Read gene expression data OmicsExpressionProteinCodingGenesTPMLogp1BatchCorrected\n",
    "try:\n",
    "    df_gene_express = pl.read_csv(\"C:\\\\Users\\\\chris\\\\rank-rx\\\\data\\\\OmicsExpressionProteinCodingGenesTPMLogp1BatchCorrected.csv\")\n",
    "    df_gene_express = df_gene_express.rename({'': 'ARXSPAN_ID'})\n",
    "    \n",
    "    # TODO when this section becomes a method, then all genes should be selected by default if no specific genes or feature selection method are declared.\n",
    "    # gene_col_names = df_gene_express.columns\n",
    "    # gene_col_names.remove('ARXSPAN_ID')\n",
    "    # Based on this research https://bmcbioinformatics.biomedcentral.com/articles/10.1186/s12859-022-04678-y\n",
    "    gene_col_names = [\n",
    "        'ZMIZ1 (57178)', 'ENG (2022)', 'FGFR1 (2260)', \n",
    "        'PAWR (5074)', 'KRT17 (3872)', 'MPO (4353)', \n",
    "        'LAT2 (7462)'\n",
    "    ]\n",
    "    \n",
    "    # # print(gene_col_names)\n",
    "    print(\"Shape of df_gene_express = {}\".format(df_gene_express.shape))\n",
    "    print(df_gene_express.head())\n",
    "except pl.errors.ParserError as e:\n",
    "    print(f\"Error: {e}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "c2080a2d-0673-485f-951d-bc438b877502",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Shape of df_dose_model = (115502, 8)\n",
      "shape: (5, 8)\n",
      "┌─────────┬────────────┬─────────────┬─────────────┬────────────┬────────┬────────────┬────────────┐\n",
      "│ DRUG_ID ┆ ARXSPAN_ID ┆ IC50_PUBLIS ┆ OncotreeCod ┆ AgeCategor ┆ Sex    ┆ PatientRac ┆ PrimaryOrM │\n",
      "│ ---     ┆ ---        ┆ HED         ┆ e           ┆ y          ┆ ---    ┆ e          ┆ etastasis  │\n",
      "│ i64     ┆ str        ┆ ---         ┆ ---         ┆ ---        ┆ str    ┆ ---        ┆ ---        │\n",
      "│         ┆            ┆ f64         ┆ str         ┆ str        ┆        ┆ str        ┆ str        │\n",
      "╞═════════╪════════════╪═════════════╪═════════════╪════════════╪════════╪════════════╪════════════╡\n",
      "│ 1003    ┆ ACH-000958 ┆ 0.025129    ┆ COAD        ┆ Adult      ┆ Female ┆ caucasian  ┆ Primary    │\n",
      "│ 1003    ┆ ACH-000651 ┆ 0.049577    ┆ COAD        ┆ Adult      ┆ Male   ┆ caucasian  ┆ Metastatic │\n",
      "│ 1003    ┆ ACH-000856 ┆ 0.028549    ┆ BRCA        ┆ Adult      ┆ Female ┆ null       ┆ Metastatic │\n",
      "│ 1003    ┆ ACH-000360 ┆ 0.039996    ┆ COAD        ┆ Adult      ┆ Male   ┆ caucasian  ┆ Metastatic │\n",
      "│ 1003    ┆ ACH-001199 ┆ 1.986678    ┆ COAD        ┆ Adult      ┆ Female ┆ asian      ┆ Primary    │\n",
      "└─────────┴────────────┴─────────────┴─────────────┴────────────┴────────┴────────────┴────────────┘\n"
     ]
    }
   ],
   "source": [
    "# Join ic50 dataset with model dataset\n",
    "# This should have been 118908 rows just like df_dose_resp_gdsc2 but it's 115502 probably because model df doesn't have some cell lines of ic50 df.\n",
    "try:\n",
    "    assert df_dose_resp_gdsc2_edited[\"ARXSPAN_ID\"].dtype == df_depmap_model_edited[\"ModelID\"].dtype\n",
    "    df_dose_model = df_dose_resp_gdsc2_edited.join(other = df_depmap_model_edited, left_on=\"ARXSPAN_ID\", right_on=\"ModelID\")\n",
    "    print(\"Shape of df_dose_model = {}\".format(df_dose_model.shape))\n",
    "    print(df_dose_model.head())\n",
    "except pl.errors.ParserError as e:\n",
    "    print(f\"Error: {e}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "1909e094-9d67-4f79-9659-86258ba75daa",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Shape of df_dose_model_gene_express = (88976, 19145)\n",
      "shape: (5, 19_145)\n",
      "┌─────────┬────────────┬───────────┬───────────┬───┬───────────┬───────────┬───────────┬───────────┐\n",
      "│ DRUG_ID ┆ ARXSPAN_ID ┆ IC50_PUBL ┆ OncotreeC ┆ … ┆ DNMT3B    ┆ ZCCHC10   ┆ PRSS2     ┆ ADAMTSL4  │\n",
      "│ ---     ┆ ---        ┆ ISHED     ┆ ode       ┆   ┆ (1789)    ┆ (54819)   ┆ (5645)    ┆ (54507)   │\n",
      "│ i64     ┆ str        ┆ ---       ┆ ---       ┆   ┆ ---       ┆ ---       ┆ ---       ┆ ---       │\n",
      "│         ┆            ┆ f64       ┆ str       ┆   ┆ f64       ┆ f64       ┆ f64       ┆ f64       │\n",
      "╞═════════╪════════════╪═══════════╪═══════════╪═══╪═══════════╪═══════════╪═══════════╪═══════════╡\n",
      "│ 1003    ┆ ACH-000958 ┆ 0.025129  ┆ COAD      ┆ … ┆ 2.072003  ┆ 4.69381   ┆ 0.240879  ┆ 0.560533  │\n",
      "│ 1003    ┆ ACH-000651 ┆ 0.049577  ┆ COAD      ┆ … ┆ 2.745968  ┆ 4.884648  ┆ 3.368963  ┆ 0.335325  │\n",
      "│ 1003    ┆ ACH-000856 ┆ 0.028549  ┆ BRCA      ┆ … ┆ 3.51348   ┆ 4.950103  ┆ -0.054202 ┆ 2.401006  │\n",
      "│ 1003    ┆ ACH-000360 ┆ 0.039996  ┆ COAD      ┆ … ┆ 0.986766  ┆ 3.83695   ┆ 0.107339  ┆ 0.082209  │\n",
      "│ 1003    ┆ ACH-000400 ┆ 0.00482   ┆ COAD      ┆ … ┆ 1.551088  ┆ 4.322617  ┆ 0.068637  ┆ 1.521804  │\n",
      "└─────────┴────────────┴───────────┴───────────┴───┴───────────┴───────────┴───────────┴───────────┘\n"
     ]
    }
   ],
   "source": [
    "# Join ic50 & model dataset with gene expression dataset\n",
    "try:\n",
    "    assert df_dose_model[\"ARXSPAN_ID\"].dtype == df_gene_express[\"ARXSPAN_ID\"].dtype\n",
    "    df_dose_model_gene_express = df_dose_model.join(df_gene_express, left_on=\"ARXSPAN_ID\", right_on=\"ARXSPAN_ID\")\n",
    "    # df_dose_model_gene_express = df_dose_model.join(df_gene_express.select(['ARXSPAN_ID'] + gene_col_names), left_on=\"ARXSPAN_ID\", right_on=\"ARXSPAN_ID\")\n",
    "    print(\"Shape of df_dose_model_gene_express = {}\".format(df_dose_model_gene_express.shape))\n",
    "    print(df_dose_model_gene_express.head())\n",
    "except Exception as e:\n",
    "    print(f\"Error: {e}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "6842a7bc-84e2-400f-aebb-8d9cfb770138",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Encode categorical features\n",
    "label_encoders = {}\n",
    "for column in ['ARXSPAN_ID', 'OncotreeCode', 'AgeCategory', 'Sex', 'PatientRace', 'PrimaryOrMetastasis']:\n",
    "    le = LabelEncoder()\n",
    "    df_dose_model_gene_express = df_dose_model_gene_express.with_columns(pl.Series(column, le.fit_transform(df_dose_model_gene_express[column].to_list())))\n",
    "    label_encoders[column] = le"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "ef3e8b03-8fdd-4cc2-b13c-ce1f827cd5d2",
   "metadata": {},
   "outputs": [],
   "source": [
    "# # Feature selection\n",
    "# pandas_df = df_dose_model_gene_express.to_pandas()\n",
    "# X_feature_selection = pandas_df.drop(columns=[target])\n",
    "# y_feature_selection = pandas_df[target]   \n",
    "# gene_col_names = mrmr_regression(X=X_feature_selection, y=y_feature_selection, K=6)\n",
    "# print(gene_col_names)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "id": "ea04a8e2-bcaf-4a9e-8470-5bf79a2438ab",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "shape: (5, 19_145)\n",
      "┌─────────┬────────────┬───────────┬───────────┬───┬───────────┬───────────┬───────────┬───────────┐\n",
      "│ DRUG_ID ┆ ARXSPAN_ID ┆ IC50_PUBL ┆ OncotreeC ┆ … ┆ DNMT3B    ┆ ZCCHC10   ┆ PRSS2     ┆ ADAMTSL4  │\n",
      "│ ---     ┆ ---        ┆ ISHED     ┆ ode       ┆   ┆ (1789)    ┆ (54819)   ┆ (5645)    ┆ (54507)   │\n",
      "│ i64     ┆ i64        ┆ ---       ┆ ---       ┆   ┆ ---       ┆ ---       ┆ ---       ┆ ---       │\n",
      "│         ┆            ┆ f64       ┆ i32       ┆   ┆ f64       ┆ f64       ┆ f64       ┆ f64       │\n",
      "╞═════════╪════════════╪═══════════╪═══════════╪═══╪═══════════╪═══════════╪═══════════╪═══════════╡\n",
      "│ 1003    ┆ 514        ┆ 0.025129  ┆ 21        ┆ … ┆ 2.072003  ┆ 4.69381   ┆ 0.240879  ┆ 0.560533  │\n",
      "│ 1003    ┆ 333        ┆ 0.049577  ┆ 21        ┆ … ┆ 2.745968  ┆ 4.884648  ┆ 3.368963  ┆ 0.335325  │\n",
      "│ 1003    ┆ 448        ┆ 0.028549  ┆ 13        ┆ … ┆ 3.51348   ┆ 4.950103  ┆ -0.054202 ┆ 2.401006  │\n",
      "│ 1003    ┆ 185        ┆ 0.039996  ┆ 21        ┆ … ┆ 0.986766  ┆ 3.83695   ┆ 0.107339  ┆ 0.082209  │\n",
      "│ 1003    ┆ 207        ┆ 0.00482   ┆ 21        ┆ … ┆ 1.551088  ┆ 4.322617  ┆ 0.068637  ┆ 1.521804  │\n",
      "└─────────┴────────────┴───────────┴───────────┴───┴───────────┴───────────┴───────────┴───────────┘\n"
     ]
    }
   ],
   "source": [
    "# Scaling numerical features\n",
    "# df_dose_model_gene_express[gene_col_names] = df_dose_model_gene_express[gene_col_names].select((pl.all()-pl.all().min()) / (pl.all().max()-pl.all().min()))\n",
    "# print(df_dose_model_gene_express.head())\n",
    "\n",
    "# Define a function to scale the columns\n",
    "def scale_column(column):\n",
    "    return (column - column.min()) / (column.max() - column.min())\n",
    "\n",
    "# Identify the columns with f64 data type\n",
    "# f64_columns = df_dose_model_gene_express.select(pl.col(gene_col_names).f64().names())\n",
    "# gene_col_names\n",
    "\n",
    "# Apply scaling to the f64 columns\n",
    "df_dose_model_gene_express = df_dose_model_gene_express.with_columns([\n",
    "    ((pl.col(col_name) - pl.col(col_name).min()) / (pl.col(col_name).max() - pl.col(col_name).min())).alias(col_name) \n",
    "    for col_name in gene_col_names\n",
    "])\n",
    "\n",
    "# Print the first few rows of the DataFrame\n",
    "print(df_dose_model_gene_express.head())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "id": "0b1906f9-97da-470c-9f50-36d1628da347",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Prepare the features\n",
    "features = ['DRUG_ID', 'ARXSPAN_ID', 'OncotreeCode', 'AgeCategory', 'Sex', 'PatientRace', 'PrimaryOrMetastasis']\n",
    "features.extend(gene_col_names)\n",
    "\n",
    "# features_grid = ['DRUG_ID', 'OncotreeCode', 'AgeCategory', 'Sex', 'PatientRace', 'PrimaryOrMetastasis']\n",
    "# features_grid.extend(gene_col_names)\n",
    "\n",
    "\n",
    "# prediction_features = ['OncotreeCode', 'AgeCategory', 'Sex', 'PatientRace', 'PrimaryOrMetastasis']\n",
    "\n",
    "# X = df_dose_model_gene_express[features]\n",
    "# print(X.head(20))\n",
    "\n",
    "\n",
    "df_dose_model_gene_express = df_dose_model_gene_express.with_columns(pl.col(target).round().cast(pl.Int32))\n",
    "\n",
    "# Scale the IC50_PUBLISHED values to 0-31 using expression\n",
    "df_dose_model_gene_express = df_dose_model_gene_express.with_columns([\n",
    "    ((pl.col(target) - pl.col(target).min()) /\n",
    "     (pl.col(target).max() - pl.col(target).min()) * 31).round().cast(pl.Int32).alias(target)\n",
    "])\n",
    "# y = df_dose_model_gene_express['IC50_PUBLISHED']\n",
    "# print(y.head(20))\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "id": "948deec4-21f5-4fe5-abd5-7d7aa645b31e",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Shape of train_data = (66551, 19145)\n",
      "shape: (20, 19_145)\n",
      "┌─────────┬────────────┬───────────┬───────────┬───┬───────────┬───────────┬───────────┬───────────┐\n",
      "│ DRUG_ID ┆ ARXSPAN_ID ┆ IC50_PUBL ┆ OncotreeC ┆ … ┆ DNMT3B    ┆ ZCCHC10   ┆ PRSS2     ┆ ADAMTSL4  │\n",
      "│ ---     ┆ ---        ┆ ISHED     ┆ ode       ┆   ┆ (1789)    ┆ (54819)   ┆ (5645)    ┆ (54507)   │\n",
      "│ i64     ┆ i64        ┆ ---       ┆ ---       ┆   ┆ ---       ┆ ---       ┆ ---       ┆ ---       │\n",
      "│         ┆            ┆ i32       ┆ i32       ┆   ┆ f64       ┆ f64       ┆ f64       ┆ f64       │\n",
      "╞═════════╪════════════╪═══════════╪═══════════╪═══╪═══════════╪═══════════╪═══════════╪═══════════╡\n",
      "│ 1003    ┆ 514        ┆ 0         ┆ 21        ┆ … ┆ 2.072003  ┆ 4.69381   ┆ 0.240879  ┆ 0.560533  │\n",
      "│ 1003    ┆ 333        ┆ 0         ┆ 21        ┆ … ┆ 2.745968  ┆ 4.884648  ┆ 3.368963  ┆ 0.335325  │\n",
      "│ 1003    ┆ 448        ┆ 0         ┆ 13        ┆ … ┆ 3.51348   ┆ 4.950103  ┆ -0.054202 ┆ 2.401006  │\n",
      "│ 1003    ┆ 185        ┆ 0         ┆ 21        ┆ … ┆ 0.986766  ┆ 3.83695   ┆ 0.107339  ┆ 0.082209  │\n",
      "│ 1003    ┆ 207        ┆ 0         ┆ 21        ┆ … ┆ 1.551088  ┆ 4.322617  ┆ 0.068637  ┆ 1.521804  │\n",
      "│ …       ┆ …          ┆ …         ┆ …         ┆ … ┆ …         ┆ …         ┆ …         ┆ …         │\n",
      "│ 1003    ┆ 492        ┆ 0         ┆ 49        ┆ … ┆ 3.755946  ┆ 4.333447  ┆ 0.015345  ┆ 0.407563  │\n",
      "│ 1003    ┆ 444        ┆ 0         ┆ 13        ┆ … ┆ 2.136818  ┆ 4.665617  ┆ 0.384091  ┆ 2.78441   │\n",
      "│ 1003    ┆ 435        ┆ 0         ┆ 51        ┆ … ┆ 2.603331  ┆ 4.978532  ┆ 0.240879  ┆ 1.458791  │\n",
      "│ 1003    ┆ 415        ┆ 0         ┆ 21        ┆ … ┆ 0.360424  ┆ 4.355636  ┆ 8.164693  ┆ -0.048902 │\n",
      "│ 1003    ┆ 4          ┆ 0         ┆ 21        ┆ … ┆ 0.407153  ┆ 5.027517  ┆ 5.548876  ┆ 0.314022  │\n",
      "└─────────┴────────────┴───────────┴───────────┴───┴───────────┴───────────┴───────────┴───────────┘\n",
      "Shape of test_data = (22425, 19145)\n",
      "shape: (20, 19_145)\n",
      "┌─────────┬────────────┬───────────┬───────────┬───┬───────────┬───────────┬───────────┬───────────┐\n",
      "│ DRUG_ID ┆ ARXSPAN_ID ┆ IC50_PUBL ┆ OncotreeC ┆ … ┆ DNMT3B    ┆ ZCCHC10   ┆ PRSS2     ┆ ADAMTSL4  │\n",
      "│ ---     ┆ ---        ┆ ISHED     ┆ ode       ┆   ┆ (1789)    ┆ (54819)   ┆ (5645)    ┆ (54507)   │\n",
      "│ i64     ┆ i64        ┆ ---       ┆ ---       ┆   ┆ ---       ┆ ---       ┆ ---       ┆ ---       │\n",
      "│         ┆            ┆ i32       ┆ i32       ┆   ┆ f64       ┆ f64       ┆ f64       ┆ f64       │\n",
      "╞═════════╪════════════╪═══════════╪═══════════╪═══╪═══════════╪═══════════╪═══════════╪═══════════╡\n",
      "│ 1003    ┆ 536        ┆ 0         ┆ 58        ┆ … ┆ 2.013664  ┆ 4.716112  ┆ 1.61707   ┆ -0.062675 │\n",
      "│ 1003    ┆ 405        ┆ 0         ┆ 51        ┆ … ┆ 4.526458  ┆ 5.909438  ┆ 0.169594  ┆ 1.050901  │\n",
      "│ 1003    ┆ 368        ┆ 0         ┆ 49        ┆ … ┆ 2.534294  ┆ 4.230593  ┆ -0.054202 ┆ 7.865699  │\n",
      "│ 1003    ┆ 208        ┆ 0         ┆ 21        ┆ … ┆ 1.454038  ┆ 4.824047  ┆ 0.015345  ┆ 1.408427  │\n",
      "│ 1003    ┆ 9          ┆ 0         ┆ 13        ┆ … ┆ 2.959773  ┆ 5.029463  ┆ 0.145017  ┆ 0.779654  │\n",
      "│ …       ┆ …          ┆ …         ┆ …         ┆ … ┆ …         ┆ …         ┆ …         ┆ …         │\n",
      "│ 1003    ┆ 540        ┆ 0         ┆ 21        ┆ … ┆ 1.379263  ┆ 5.07402   ┆ 1.558514  ┆ -0.035256 │\n",
      "│ 1003    ┆ 6          ┆ 0         ┆ 21        ┆ … ┆ 5.933968  ┆ 5.327266  ┆ 1.84497   ┆ 3.581425  │\n",
      "│ 1003    ┆ 176        ┆ 0         ┆ 21        ┆ … ┆ 2.010529  ┆ 4.653817  ┆ 0.205682  ┆ 5.227169  │\n",
      "│ 1003    ┆ 78         ┆ 0         ┆ 54        ┆ … ┆ 0.736067  ┆ 4.134357  ┆ 0.205682  ┆ 2.377232  │\n",
      "│ 1003    ┆ 399        ┆ 0         ┆ 57        ┆ … ┆ 2.5365    ┆ 4.687726  ┆ 0.015345  ┆ 0.841034  │\n",
      "└─────────┴────────────┴───────────┴───────────┴───┴───────────┴───────────┴───────────┴───────────┘\n",
      "Shape of X_train = (66551, 14)\n",
      "shape: (20, 14)\n",
      "┌─────────┬────────────┬───────────┬───────────┬───┬───────────┬───────────┬───────────┬───────────┐\n",
      "│ DRUG_ID ┆ ARXSPAN_ID ┆ OncotreeC ┆ AgeCatego ┆ … ┆ PAWR      ┆ KRT17     ┆ MPO       ┆ LAT2      │\n",
      "│ ---     ┆ ---        ┆ ode       ┆ ry        ┆   ┆ (5074)    ┆ (3872)    ┆ (4353)    ┆ (7462)    │\n",
      "│ i64     ┆ i64        ┆ ---       ┆ ---       ┆   ┆ ---       ┆ ---       ┆ ---       ┆ ---       │\n",
      "│         ┆            ┆ i32       ┆ i64       ┆   ┆ f64       ┆ f64       ┆ f64       ┆ f64       │\n",
      "╞═════════╪════════════╪═══════════╪═══════════╪═══╪═══════════╪═══════════╪═══════════╪═══════════╡\n",
      "│ 1003    ┆ 514        ┆ 21        ┆ 0         ┆ … ┆ 0.679258  ┆ 0.192601  ┆ 0.110178  ┆ 0.183139  │\n",
      "│ 1003    ┆ 333        ┆ 21        ┆ 0         ┆ … ┆ 0.607023  ┆ 0.109392  ┆ 0.001061  ┆ 0.448888  │\n",
      "│ 1003    ┆ 448        ┆ 13        ┆ 0         ┆ … ┆ 0.714077  ┆ 0.108544  ┆ 0.041807  ┆ 0.087857  │\n",
      "│ 1003    ┆ 185        ┆ 21        ┆ 0         ┆ … ┆ 0.547431  ┆ 0.08301   ┆ 0.039623  ┆ 0.062555  │\n",
      "│ 1003    ┆ 207        ┆ 21        ┆ 0         ┆ … ┆ 0.744917  ┆ 0.33121   ┆ 0.0       ┆ 0.029978  │\n",
      "│ …       ┆ …          ┆ …         ┆ …         ┆ … ┆ …         ┆ …         ┆ …         ┆ …         │\n",
      "│ 1003    ┆ 492        ┆ 49        ┆ 0         ┆ … ┆ 0.768595  ┆ 0.049699  ┆ 0.027978  ┆ 0.061635  │\n",
      "│ 1003    ┆ 444        ┆ 13        ┆ 0         ┆ … ┆ 0.67151   ┆ 0.76532   ┆ 0.043947  ┆ 0.190074  │\n",
      "│ 1003    ┆ 435        ┆ 51        ┆ 0         ┆ … ┆ 0.67733   ┆ 0.049699  ┆ 0.0       ┆ 0.106842  │\n",
      "│ 1003    ┆ 415        ┆ 21        ┆ 2         ┆ … ┆ 0.602568  ┆ 0.489361  ┆ 0.038885  ┆ 0.175437  │\n",
      "│ 1003    ┆ 4          ┆ 21        ┆ 0         ┆ … ┆ 0.633101  ┆ 0.06601   ┆ 0.016743  ┆ 0.109565  │\n",
      "└─────────┴────────────┴───────────┴───────────┴───┴───────────┴───────────┴───────────┴───────────┘\n",
      "Shape of y_train = (66551, 1)\n",
      "shape: (20, 1)\n",
      "┌────────────────┐\n",
      "│ IC50_PUBLISHED │\n",
      "│ ---            │\n",
      "│ i32            │\n",
      "╞════════════════╡\n",
      "│ 0              │\n",
      "│ 0              │\n",
      "│ 0              │\n",
      "│ 0              │\n",
      "│ 0              │\n",
      "│ …              │\n",
      "│ 0              │\n",
      "│ 0              │\n",
      "│ 0              │\n",
      "│ 0              │\n",
      "│ 0              │\n",
      "└────────────────┘\n",
      "Shape of X_test = (22425, 14)\n",
      "shape: (20, 14)\n",
      "┌─────────┬────────────┬───────────┬───────────┬───┬───────────┬───────────┬───────────┬───────────┐\n",
      "│ DRUG_ID ┆ ARXSPAN_ID ┆ OncotreeC ┆ AgeCatego ┆ … ┆ PAWR      ┆ KRT17     ┆ MPO       ┆ LAT2      │\n",
      "│ ---     ┆ ---        ┆ ode       ┆ ry        ┆   ┆ (5074)    ┆ (3872)    ┆ (4353)    ┆ (7462)    │\n",
      "│ i64     ┆ i64        ┆ ---       ┆ ---       ┆   ┆ ---       ┆ ---       ┆ ---       ┆ ---       │\n",
      "│         ┆            ┆ i32       ┆ i64       ┆   ┆ f64       ┆ f64       ┆ f64       ┆ f64       │\n",
      "╞═════════╪════════════╪═══════════╪═══════════╪═══╪═══════════╪═══════════╪═══════════╪═══════════╡\n",
      "│ 1003    ┆ 536        ┆ 58        ┆ 0         ┆ … ┆ 0.808868  ┆ 0.040992  ┆ 0.001061  ┆ 0.058842  │\n",
      "│ 1003    ┆ 405        ┆ 51        ┆ 0         ┆ … ┆ 0.861952  ┆ 0.113132  ┆ 0.027155  ┆ 0.04302   │\n",
      "│ 1003    ┆ 368        ┆ 49        ┆ 0         ┆ … ┆ 0.53094   ┆ 0.037165  ┆ 0.010164  ┆ 0.391965  │\n",
      "│ 1003    ┆ 208        ┆ 21        ┆ 0         ┆ … ┆ 0.772571  ┆ 0.14041   ┆ 0.01855   ┆ 0.249758  │\n",
      "│ 1003    ┆ 9          ┆ 13        ┆ 0         ┆ … ┆ 0.690253  ┆ 0.134415  ┆ 0.0       ┆ 0.039862  │\n",
      "│ …       ┆ …          ┆ …         ┆ …         ┆ … ┆ …         ┆ …         ┆ …         ┆ …         │\n",
      "│ 1003    ┆ 540        ┆ 21        ┆ 0         ┆ … ┆ 0.56923   ┆ 0.202225  ┆ 0.022939  ┆ 0.068844  │\n",
      "│ 1003    ┆ 6          ┆ 21        ┆ 0         ┆ … ┆ 0.750837  ┆ 0.239278  ┆ 0.001061  ┆ 0.040922  │\n",
      "│ 1003    ┆ 176        ┆ 21        ┆ 0         ┆ … ┆ 0.731969  ┆ 0.168084  ┆ 0.03279   ┆ 0.166864  │\n",
      "│ 1003    ┆ 78         ┆ 54        ┆ 2         ┆ … ┆ 0.705942  ┆ 0.088173  ┆ 0.025489  ┆ 0.196325  │\n",
      "│ 1003    ┆ 399        ┆ 57        ┆ 0         ┆ … ┆ 0.863131  ┆ 0.059109  ┆ 0.04878   ┆ 0.0579    │\n",
      "└─────────┴────────────┴───────────┴───────────┴───┴───────────┴───────────┴───────────┴───────────┘\n",
      "Shape of y_test = (22425, 1)\n",
      "shape: (20, 1)\n",
      "┌────────────────┐\n",
      "│ IC50_PUBLISHED │\n",
      "│ ---            │\n",
      "│ i32            │\n",
      "╞════════════════╡\n",
      "│ 0              │\n",
      "│ 0              │\n",
      "│ 0              │\n",
      "│ 0              │\n",
      "│ 0              │\n",
      "│ …              │\n",
      "│ 0              │\n",
      "│ 0              │\n",
      "│ 0              │\n",
      "│ 0              │\n",
      "│ 0              │\n",
      "└────────────────┘\n"
     ]
    }
   ],
   "source": [
    "# Group Shuffle and Split using Polars\n",
    "def group_shuffle_split(df, group_col, test_size=0.2, random_state=42):\n",
    "    np.random.seed(random_state)\n",
    "    groups = df[group_col].unique().to_list()\n",
    "    np.random.shuffle(groups)\n",
    "    test_groups = groups[:int(test_size * len(groups))]\n",
    "    train_groups = groups[int(test_size * len(groups)):]\n",
    "    \n",
    "    train_df = df.filter(pl.col(group_col).is_in(train_groups))\n",
    "    test_df = df.filter(pl.col(group_col).is_in(test_groups))\n",
    "    \n",
    "    return train_df, test_df\n",
    "\n",
    "train_data, test_data = group_shuffle_split(df_dose_model_gene_express, 'ARXSPAN_ID', test_size=0.25, random_state=42)\n",
    "print(\"Shape of train_data = {}\".format(train_data.shape))\n",
    "print(train_data.head(20))\n",
    "print(\"Shape of test_data = {}\".format(test_data.shape))\n",
    "print(test_data.head(20))\n",
    "\n",
    "\n",
    "# Separate features and target in train and test data\n",
    "X_train = train_data.select(features)\n",
    "print(\"Shape of X_train = {}\".format(X_train.shape))\n",
    "print(X_train.head(20))\n",
    "\n",
    "y_train = train_data.select([target])\n",
    "print(\"Shape of y_train = {}\".format(y_train.shape))\n",
    "print(y_train.head(20))\n",
    "\n",
    "X_test = test_data.select(features)\n",
    "print(\"Shape of X_test = {}\".format(X_test.shape))\n",
    "print(X_test.head(20))\n",
    "\n",
    "y_test = test_data.select([target])\n",
    "print(\"Shape of y_test = {}\".format(y_test.shape))\n",
    "print(y_test.head(20))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "id": "070f4b7d-74f2-4a4c-9f67-b7203c8e1324",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "shape of group_train: 452\n",
      "shape of group_test: 150\n",
      "Shape of dtrain DMatrix: (66551, 14)\n",
      "Shape of dtest DMatrix: (22425, 14)\n"
     ]
    }
   ],
   "source": [
    "# Create group parameter for XGBoost\n",
    "group_train = X_train.group_by('ARXSPAN_ID').count().select('count').to_series().to_list()\n",
    "print(f\"shape of group_train: {len(group_train)}\")\n",
    "group_test = X_test.group_by('ARXSPAN_ID').count().select('count').to_series().to_list()\n",
    "print(f\"shape of group_test: {len(group_test)}\")\n",
    "\n",
    "# Convert data to DMatrix\n",
    "dtrain = xgb.DMatrix(X_train.select(features).to_numpy(), label=y_train.to_numpy())\n",
    "dtrain.set_group(group_train)\n",
    "print(f\"Shape of dtrain DMatrix: ({dtrain.num_row()}, {dtrain.num_col()})\")\n",
    "\n",
    "dtest = xgb.DMatrix(X_test.select(features).to_numpy(), label=y_test.to_numpy())\n",
    "dtest.set_group(group_test)\n",
    "print(f\"Shape of dtest DMatrix: ({dtest.num_row()}, {dtest.num_col()})\")\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "id": "c8eca6d3-f15a-41ff-a4eb-16423a65c8cc",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Define XGBoost parameters\n",
    "params = {\n",
    "    'objective': 'rank:pairwise',\n",
    "    'eta': 0.1,\n",
    "    'gamma': 1.0,\n",
    "    'min_child_weight': 0.5,\n",
    "    'max_depth': 5,\n",
    "    'eval_metric': 'ndcg'\n",
    "}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "id": "b523cba1-f7f1-446e-a4da-f2bc45143267",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[0]\ttest-ndcg:0.23861\n",
      "[1]\ttest-ndcg:0.27940\n",
      "[2]\ttest-ndcg:0.29439\n",
      "[3]\ttest-ndcg:0.28680\n",
      "[4]\ttest-ndcg:0.29420\n",
      "[5]\ttest-ndcg:0.29220\n",
      "[6]\ttest-ndcg:0.29884\n",
      "[7]\ttest-ndcg:0.28782\n",
      "[8]\ttest-ndcg:0.28554\n",
      "[9]\ttest-ndcg:0.30076\n",
      "[10]\ttest-ndcg:0.30157\n",
      "[11]\ttest-ndcg:0.30191\n",
      "[12]\ttest-ndcg:0.30574\n",
      "[13]\ttest-ndcg:0.31183\n",
      "[14]\ttest-ndcg:0.31187\n",
      "[15]\ttest-ndcg:0.31310\n",
      "[16]\ttest-ndcg:0.33289\n",
      "[17]\ttest-ndcg:0.33825\n",
      "[18]\ttest-ndcg:0.33373\n",
      "[19]\ttest-ndcg:0.34136\n",
      "[20]\ttest-ndcg:0.35748\n",
      "[21]\ttest-ndcg:0.35686\n",
      "[22]\ttest-ndcg:0.35686\n",
      "[23]\ttest-ndcg:0.35865\n",
      "[24]\ttest-ndcg:0.36701\n",
      "[25]\ttest-ndcg:0.36716\n",
      "[26]\ttest-ndcg:0.36519\n",
      "[27]\ttest-ndcg:0.36520\n",
      "[28]\ttest-ndcg:0.35742\n",
      "[29]\ttest-ndcg:0.35742\n",
      "[30]\ttest-ndcg:0.35619\n",
      "[31]\ttest-ndcg:0.35069\n",
      "[32]\ttest-ndcg:0.35801\n",
      "[33]\ttest-ndcg:0.34271\n",
      "[34]\ttest-ndcg:0.36531\n",
      "[35]\ttest-ndcg:0.36353\n"
     ]
    }
   ],
   "source": [
    "# Train the model\n",
    "model = xgb.train(params, dtrain, num_boost_round=100, evals=[(dtest, 'test')], early_stopping_rounds=10)\n",
    "\n",
    "# Predict and evaluate\n",
    "y_pred = model.predict(dtest)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "id": "ca704930-77cb-473c-bd3f-4588e3484c83",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "shape: (602, 2)\n",
      "┌────────────┬─────┐\n",
      "│ ARXSPAN_ID ┆ len │\n",
      "│ ---        ┆ --- │\n",
      "│ i64        ┆ u32 │\n",
      "╞════════════╪═════╡\n",
      "│ 375        ┆ 12  │\n",
      "│ 381        ┆ 22  │\n",
      "│ 185        ┆ 25  │\n",
      "│ 108        ┆ 29  │\n",
      "│ 432        ┆ 47  │\n",
      "│ …          ┆ …   │\n",
      "│ 398        ┆ 173 │\n",
      "│ 139        ┆ 173 │\n",
      "│ 294        ┆ 173 │\n",
      "│ 163        ┆ 173 │\n",
      "│ 315        ┆ 173 │\n",
      "└────────────┴─────┘\n"
     ]
    }
   ],
   "source": [
    "cell_line_groups = df_dose_model_gene_express.group_by('ARXSPAN_ID').len()\n",
    "print(cell_line_groups.sort(by='len'))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e4c17e17-e3df-4b26-82ca-dd86d200dd64",
   "metadata": {},
   "source": []
  }
 ],
 "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.10.14"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}