Train_FLAMES.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "218984fe",
   "metadata": {},
   "source": [
    "# Import packages"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 30,
   "id": "e4d8ba17",
   "metadata": {},
   "outputs": [],
   "source": [
    "import pandas as pd\n",
    "import numpy as np\n",
    "from sklearn.model_selection import train_test_split, RandomizedSearchCV\n",
    "import xgboost as xgb\n",
    "import random\n",
    "import os\n",
    "from sklearn.metrics import auc, precision_recall_curve\n",
    "from sklearn.preprocessing import MinMaxScaler\n",
    "import pickle"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d5da8c8e",
   "metadata": {},
   "source": [
    "# Define functions"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 31,
   "id": "b7553c64",
   "metadata": {},
   "outputs": [],
   "source": [
    "def create_relative_annotation(genes):\n",
    "    for feature in genes.columns:\n",
    "        if len(genes) < 1:\n",
    "            continue\n",
    "        if \"rel_\" in feature:\n",
    "            continue\n",
    "        if \"distance\" in feature:\n",
    "            maxval = max(genes[feature].astype(float))\n",
    "            minval = min(genes[feature].astype(float))\n",
    "            if maxval == minval:\n",
    "                genes[\"rel_\" + feature] = 1\n",
    "            else:\n",
    "                genes[\"rel_\" + feature] = maxval - genes[feature]\n",
    "                genes[\"rel_\" + feature] = genes[\"rel_\" + feature] / float(\n",
    "                    maxval - minval\n",
    "                )\n",
    "        elif any(x in feature for x in [\"max\", \"sum\", \"PoPS\", \"MAGMA_Z\"]):\n",
    "            maxval = max(genes[feature].astype(float))\n",
    "            minval = min(genes[feature].astype(float))\n",
    "            if maxval == minval and abs(maxval) > 0.0:\n",
    "                genes[\"rel_\" + feature] = 1\n",
    "            elif maxval == 0.0:\n",
    "                genes[f\"rel_{feature}\"] = 0\n",
    "            else:\n",
    "                genes[f\"rel_{feature}\"] = (genes[feature].astype(float) - minval) / (maxval - minval)\n",
    "    return genes"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4f471947",
   "metadata": {},
   "source": [
    "# Train FLAMES (LOCO) models"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 33,
   "id": "a0752a23",
   "metadata": {},
   "outputs": [],
   "source": [
    "#load training data\n",
    "df = pd.read_csv(os.path.join('annotated_loci', 'all_FLAMES.csv'))\n",
    "df = df[df['type'] == 'protein_coding']\n",
    "df = df.fillna(0)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 34,
   "id": "d65e5db6",
   "metadata": {},
   "outputs": [],
   "source": [
    "feats = ['rel_MAGMA_Z', 'VEP_max', 'rel_VEP_max', 'CADD_sum', 'rel_CADD_sum', 'CLPP_GTEX_tissue_weighted_eQTL_sum', 'rel_CLPP_GTEX_tissue_weighted_eQTL_sum', 'RoadmapEpi_sum', 'rel_RoadmapEpi_sum', 'HACER_PRO-seq_GRO-seq_sum', 'rel_HACER_PRO-seq_GRO-seq_sum', 'Promoter_sum', 'rel_Promoter_sum', 'F5_ep_sum', 'rel_F5_ep_sum', 'Jung_HiC_sum', 'rel_Jung_HiC_sum', 'GeneHancer_sum', 'rel_GeneHancer_sum', 'EpiMap_sum', 'rel_EpiMap_sum', 'mean_CLPP_eQTL_eQTLCatalog_max', 'rel_mean_CLPP_eQTL_eQTLCatalog_max', 'ABC_CRISPR_sum', 'rel_ABC_CRISPR_sum', 'ABC_EP_sum', 'rel_ABC_EP_sum', 'max_CLPP_rQTL_eQTLCatalog_max', 'rel_max_CLPP_rQTL_eQTLCatalog_max', 'Javierre_HiC_sum', 'rel_Javierre_HiC_sum', 'HACER_CAGE_max', 'rel_HACER_CAGE_max', 'CLPP_eQTLgen_sum', 'rel_CLPP_eQTLgen_sum', 'cicero_sum', 'rel_cicero_sum', 'MAGMA_Z', 'rel_weighted_distance', 'rel_weighted_TSS_distance', 'weighted_distance']\n",
    "df = df.groupby('filename').apply(create_relative_annotation)\n",
    "df= df.reset_index(drop=True)\n",
    "traindf = df.copy()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "382b7384",
   "metadata": {},
   "outputs": [],
   "source": [
    "#Train LOCO models\n",
    "LOCO_models = []\n",
    "sav_params = []\n",
    "iters = []\n",
    "for chrom in range(1,23):    \n",
    "    keeps = set(df[df['chr'] != chrom]['filename'])\n",
    "    #train in LOCO framework\n",
    "    if len(keeps) < 100:\n",
    "        extra = set(df[df['chr'] != 23-chrom]['filename'])\n",
    "        if len(extra) + len(keeps) >= 100:\n",
    "            keeps.append(extra[0:100-len(keeps)])\n",
    "        else:\n",
    "            keeps.append(extra)\n",
    "    \n",
    "                  \n",
    "    X_train = df[df['filename'].isin(keeps)][feats]\n",
    "    y_train = df[df['filename'].isin(keeps)]['TP']\n",
    "    \n",
    "    X_valid = df[~ df['filename'].isin(keeps)][feats]\n",
    "    y_valid = df[~ df['filename'].isin(keeps)]['TP']\n",
    "    eval_set = [(X_train, y_train), (X_valid, y_valid)]\n",
    "    \n",
    "    # Define the parameter grid for random search\n",
    "    param_dist = {\n",
    "        \n",
    "        \"learning_rate\": [0.1],\n",
    "        'max_depth': [3, 4, 5],\n",
    "        'colsample_bytree': [0.8,0.9],\n",
    "        'min_child_weight': [3, 4, 5]\n",
    "        # Add more parameters as needed\n",
    "    }\n",
    "\n",
    "    # Instantiate the XGBoost classifier\n",
    "    clf = xgb.XGBClassifier(objective='binary:logistic', seed=42, eval_metric='rmse', early_stopping_rounds=7)\n",
    "\n",
    "    # Perform random search\n",
    "    random_search = RandomizedSearchCV(clf, param_distributions=param_dist, n_iter=10, cv=5, scoring='accuracy', random_state=42, verbose =False)\n",
    "        \n",
    "    fit_params={ \n",
    "            \"eval_set\" : [[X_valid, y_valid]],\n",
    "                \"verbose\": False\n",
    "               }\n",
    "    \n",
    "    random_search.fit(X_train, y_train, **fit_params)\n",
    "\n",
    "    # Print the best parameters found\n",
    "    print(\"Best parameters found: \", random_search.best_params_)\n",
    "\n",
    "    # Get the best estimator\n",
    "    best_model = random_search.best_estimator_\n",
    "    num_trees = best_model.best_iteration\n",
    "\n",
    "    # Fit the best model on the training data with early stopping\n",
    "    \n",
    "    best_model.fit(X_train, y_train, eval_set=eval_set, verbose=True)\n",
    "  \n",
    "    # Evaluate the best model on test data\n",
    "    LOCO_models.append(best_model)\n",
    "    sav_params.append(random_search.best_params_)\n",
    "    iters.append(num_trees)\n",
    "\n",
    "    print(f'Finnished chrom {chrom}')\n",
    "print('done')\n",
    "\n",
    "final_params = {}\n",
    "for key in sav_params[0]:\n",
    "    final_params[key] = round(np.mean([x[key] for x in sav_params]),2)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "70659822",
   "metadata": {},
   "outputs": [],
   "source": [
    "# #Train consensus\n",
    "random_locs = list(set(traindf['filename']))\n",
    "random.seed(len('FLAMES reproducibility'))\n",
    "random.shuffle(random_locs)\n",
    "holdout = random_locs[0:int(len(set(traindf['filename']))/9)]\n",
    "\n",
    "traindf = df\n",
    "filename_counts = traindf['filename'].value_counts()\n",
    "\n",
    "\n",
    "X_train = traindf[~traindf['filename'].isin(holdout)][feats]\n",
    "y_train = traindf[~traindf['filename'].isin(holdout)]['TP']\n",
    "X_valid = traindf[traindf['filename'].isin(holdout)][feats]\n",
    "y_valid = traindf[traindf['filename'].isin(holdout)]['TP']\n",
    "\n",
    "eval_set = [(X_train, y_train), (X_valid, y_valid)]\n",
    "    \n",
    "FLAMES_xgb = xgb.XGBClassifier(objective='binary:logistic', seed=42, min_child_weight=int(np.round(final_params['min_child_weight'],0)), max_depth=int(np.round(final_params['max_depth'],0)), learning_rate=final_params['learning_rate'], colsample_bytree=final_params['colsample_bytree'],  eval_metric='rmse', early_stopping_rounds=10)\n",
    "FLAMES_xgb.fit(X_train, y_train, eval_set=eval_set, verbose=True)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "30ba59b2",
   "metadata": {},
   "source": [
    "# Calibrate XGB & PoPS integration"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1863e9c2",
   "metadata": {},
   "outputs": [],
   "source": [
    "xdf = df.drop_duplicates(keep=\"first\")\n",
    "xdf = xdf.drop_duplicates(subset=[\"ensg\", \"symbol\", 'filename'], keep=\"first\")\n",
    "new_df=[]\n",
    "for chrom in range(1,23):    \n",
    "    tdf = xdf[xdf['chr'] == chrom].copy()\n",
    "    X_train = tdf[feats]\n",
    "    y = [x[1] for x in LOCO_models[chrom-1].predict_proba(X_train)]\n",
    "    tdf['XGB_score'] = y\n",
    "    new_df.append(tdf)\n",
    "results = pd.concat(new_df)\n",
    "\n",
    "print(len(set(results['filename'])), sum(results['TP']))\n",
    "best_x = []\n",
    "best_y =[]\n",
    "best = 0\n",
    "for x in np.linspace(0,0.9999999, 5):\n",
    "    print(np.round(x, 2))\n",
    "    for y in np.linspace(0,2,5):\n",
    "        test_results = results.groupby('filename').apply(Flames_scoring_light, x, y) \n",
    "        if sum(test_results['TP'] * test_results['FLAMES_causal'])/sum(test_results['FLAMES_causal']) > best:\n",
    "            best = sum(test_results['TP'] * test_results['FLAMES_causal'])/sum(test_results['FLAMES_causal'])\n",
    "            best_combi = (np.round(x,2),np.round(y,2))\n",
    "print(best_combi)\n",
    "            \n",
    "for x in np.linspace(best_combi[0] - 0.125,best_combi[0] + 0.125, 5):\n",
    "    print(np.round(x, 4))\n",
    "    for y in np.linspace(best_combi[1]-0.25,best_combi[1]+0.25,5):\n",
    "        test_results = results.groupby('filename').apply(Flames_scoring_light, x, y) \n",
    "        if sum(test_results['TP'] * test_results['FLAMES_causal'])/sum(test_results['FLAMES_causal']) > best:\n",
    "            best = sum(test_results['TP'] * test_results['FLAMES_causal'])/sum(test_results['FLAMES_causal'])\n",
    "            best_x = [x]\n",
    "            best_y = [y]\n",
    "        elif sum(test_results['TP'] * test_results['FLAMES_causal'])/sum(test_results['FLAMES_causal']) == best:\n",
    "            best = sum(test_results['TP'] * test_results['FLAMES_causal'])/sum(test_results['FLAMES_causal'])\n",
    "            best_x.append(x)\n",
    "            best_y.append(y)\n",
    "print(np.mean(best_x), np.mean(best_y))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8f6aa0aa",
   "metadata": {},
   "source": [
    "# Calibrate FLAMES framework"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 36,
   "id": "e0c8fef6",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Scaled score threshold:0.248\n",
      "Raw score threshold: 0.136\n",
      "Estimated recall at thresholds: 0.364\n"
     ]
    }
   ],
   "source": [
    "df = pd.read_csv(os.path.join('benchmarking_results', 'calibration', 'calibration.tsv'), sep='\\t')\n",
    "ts = []\n",
    "rs = []\n",
    "recs = []\n",
    "for i in range(1000):\n",
    "    best_r = 0\n",
    "    thresh = 0\n",
    "    min_raw = 0\n",
    "    sampledf = df.sample(replace=True, frac=1, random_state=i)\n",
    "    for x in np.linspace(1,0,26):\n",
    "        tempdf = sampledf[sampledf['FLAMES_causal'] == 1]\n",
    "        tempdf = tempdf[tempdf['FLAMES_raw'] > x]\n",
    "        tempdf = tempdf.sort_values('FLAMES_scaled', ascending= False)\n",
    "        if len(tempdf) == 0:\n",
    "            continue\n",
    "        cumsum = np.cumsum(tempdf['TP'])\n",
    "        rank = np.arange(len(cumsum)) + 1\n",
    "        Num = sum(sampledf['TP'])\n",
    "        precisions = cumsum / rank\n",
    "        recalls = cumsum / Num\n",
    "        last = 0\n",
    "        for i in range(len(precisions)):\n",
    "            if list(precisions)[i] >= 0.75:\n",
    "                last = i\n",
    "        i = last\n",
    "        precision_at_threshold = list(precisions)[i]\n",
    "        recall_at_threshold = list(recalls)[i]\n",
    "        if recall_at_threshold > best_r:\n",
    "            best_r = recall_at_threshold\n",
    "            thresh = list(tempdf['FLAMES_scaled'])[i]\n",
    "            min_raw = np.round(x,2)\n",
    "    ts.append(thresh)\n",
    "    rs.append(min_raw)\n",
    "    recs.append(best_r)\n",
    "print(f'Scaled score threshold:{np.round(np.mean(ts),3)}\\nRaw score threshold: {np.round(np.mean(rs),3)}\\nEstimated recall at thresholds: {np.round(np.mean(recs),3)}')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 28,
   "id": "aa785068",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Raw score threshold:0.307\n",
      "Estimated recall at thresholds: 0.273\n"
     ]
    }
   ],
   "source": [
    "df = pd.read_csv(os.path.join('benchmarking_results', 'calibration', 'calibration.tsv'), sep='\\t')\n",
    "ts = []\n",
    "recs = []\n",
    "bsdf = []\n",
    "for i in range(1000):\n",
    "    best_r = 0\n",
    "    thresh = 0\n",
    "    min_raw = 0\n",
    "    sampledf = df.sample(replace=True, frac=1, random_state=i)\n",
    "    tempdf = sampledf[sampledf['FLAMES_causal'] == 1]\n",
    "    tempdf = tempdf.sort_values('FLAMES_raw', ascending= False)\n",
    "    cumsum = np.cumsum(tempdf['TP'])\n",
    "    rank = np.arange(len(cumsum)) + 1\n",
    "    Num = sum(sampledf['TP'])\n",
    "    precisions = cumsum / rank\n",
    "    recalls = cumsum / Num\n",
    "    last = 0\n",
    "    bsdf.append(tempdf)\n",
    "    for i in range(len(precisions)):\n",
    "        if list(precisions)[i] >= 0.75:\n",
    "            last = i\n",
    "    i = last\n",
    "    precision_at_threshold = list(precisions)[i]\n",
    "    recall_at_threshold = list(recalls)[i]\n",
    "    if recall_at_threshold > best_r:\n",
    "        best_r = recall_at_threshold\n",
    "        thresh = list(tempdf['FLAMES_scaled'])[i]\n",
    "    ts.append(thresh)\n",
    "    recs.append(best_r)\n",
    "print(f'Raw score threshold:{np.round(np.mean(ts),3)}\\nEstimated recall at thresholds: {np.round(np.mean(recs),3)}')"
   ]
  }
 ],
 "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.8.13"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}