{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "d408841e",
   "metadata": {},
   "source": [
    "# ROSMAP (Alzheimer's Disease) Biomarker Discovery\n",
    "\n",
    "**Cohort Summary:**\n",
    "\n",
    "| Stage | Methylation | mRNA | miRNA | Clinical |\n",
    "| --- | --- | --- | --- | --- |\n",
    "| **Raw** (features x samples) | 420,132 x 741 | 55,889 x 642 | 309 x 704 | 18 x 3584 |\n",
    "| **Data Aligned** (samples x features) | 521 x 420,132 | 521 x 55,889 | 521 x 309 | 521 x 17 |\n",
    "| **Variance top 100k** (Methylation only) | 521 x 100,000 | - | - | - |\n",
    "| **Anova-F & Random Forest Intersection** k=8000 | 521 x 772 | 521 x 2054 | - | - |\n",
    "| **After feature selection** | 521 x 772 | 521 x 2054 | 521 x 309 | 521 x 17 |\n",
    "| **Final feature selection** (Variance k=300) | 521 x 300 | 521 x 300 | 521 x 300 | 521 x 17 |\n",
    "\n",
    "**Target Definition:** Multi-class cognitive diagnosis derived from clinical cogdx scores. No Cognitive Impairment (NCI, n=168), Mild Cognitive Impairment (MCI, n=134), and Dementia (n=219).\n",
    "\n",
    "\n",
    "**Data Source:** [AD Knowledge Portal ROSMAP](https://adknowledgeportal.synapse.org/Explore/Studies/DetailsPage/StudyData?Study=syn3219045)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "40c36cd7",
   "metadata": {},
   "outputs": [],
   "source": [
    "!pip install synapseclient pandas bioneuralnet"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0ecb6d27",
   "metadata": {},
   "outputs": [],
   "source": [
    "# torch CPU\n",
    "!pip install --no-cache-dir torch==2.5.1+cpu --index-url https://download.pytorch.org/whl/cpu\n",
    "\n",
    "# torch_geometric\n",
    "!pip install --no-cache-dir pyg-lib torch-scatter torch-sparse torch-cluster torch-spline-conv torch-geometric -f https://data.pyg.org/whl/torch-2.5.1+cpu.html\n",
    "\n",
    "# tunning for DPMON\n",
    "!pip install --no-cache-dir \"ray[tune]==2.46.0\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c5e36951",
   "metadata": {},
   "outputs": [],
   "source": [
    "import synapseclient\n",
    "\n",
    "TOKEN = \"your_synapse_token\"\n",
    "\n",
    "try:\n",
    "    syn = synapseclient.login(authToken=TOKEN)\n",
    "    print(f\"Login successful! User: {syn.getUserProfile().userName}\")\n",
    "except Exception as e:\n",
    "    print(\"Login failed. Check that you copied the full token.\")\n",
    "    raise e\n",
    "\n",
    "# you should recieve a Welcome + Login Successful message"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8a36e953",
   "metadata": {},
   "outputs": [],
   "source": [
    "import pandas as pd\n",
    "\n",
    "try:\n",
    "    syn = synapseclient.login(authToken=TOKEN)\n",
    "    print(f\"Login successful! User: {syn.getUserProfile().userName}\")\n",
    "except Exception as e:\n",
    "    print(\"Login failed. Check that you copied the full token.\")\n",
    "    raise e\n",
    "\n",
    "files_to_download = {\n",
    "    \"ROSMAP_assay_RNAseq_metadata.csv\": {\"id\": \"syn21088596\", \"ver\": 5},\n",
    "    \"ROSMAP_assay_methylationArray_metadata.csv\": {\"id\": \"syn23569437\", \"ver\": 1},\n",
    "    \"ROSMAP_assay_miRNAarray_nanostring_metadata.csv\": {\"id\": \"syn23569439\", \"ver\": 4},\n",
    "    \"ROSMAP_arrayMethylation_imputed.tsv.gz\": {\"id\": \"syn3168763\",  \"ver\": 2},\n",
    "    \"ROSMAP_clinical.csv\": {\"id\": \"syn3191087\",  \"ver\": 11},\n",
    "    \"ROSMAP_clinical_codebook.pdf\": {\"id\": \"syn3191090\",  \"ver\": 2},\n",
    "    \"ROSMAP_arraymiRNA.gct\": {\"id\": \"syn3387327\",  \"ver\": 1},\n",
    "    \"ROSMAP_RNAseq_FPKM_gene.tsv\": {\"id\": \"syn3505720\",  \"ver\": 2},\n",
    "    \"ROSMAP_biospecimen_metadata.csv\": {\"id\": \"syn21323366\", \"ver\": 1},\n",
    "    \"ROSMAP_IDkey.csv\": {\"id\": \"syn3382527\", \"ver\": 9},\n",
    "}\n",
    "\n",
    "downloaded_files = {}\n",
    "\n",
    "print(\"Starting download\")\n",
    "for filename, info in files_to_download.items():\n",
    "    print(f\"Downloading {filename}\")\n",
    "    \n",
    "    entity = syn.get(info[\"id\"], version=info[\"ver\"])\n",
    "    \n",
    "    downloaded_files[filename] = entity.path\n",
    "\n",
    "print(\"\\nAll files downloaded\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7067b333",
   "metadata": {},
   "outputs": [],
   "source": [
    "# clean and transpose the ROSMAP data\n",
    "rna_path = downloaded_files[\"ROSMAP_RNAseq_FPKM_gene.tsv\"]\n",
    "rna_meta_path = downloaded_files[\"ROSMAP_assay_RNAseq_metadata.csv\"]\n",
    "\n",
    "# get the valid specimenIDs\n",
    "df_meta = pd.read_csv(rna_meta_path)\n",
    "print(f\"Metadata Columns: {df_meta.columns.tolist()}\")\n",
    "\n",
    "# we use a set for faster lookup of valid IDs\n",
    "valid_specimens = set(df_meta['specimenID'].astype(str))\n",
    "print(f\"Found {len(valid_specimens)} valid specimenIDs in metadata\")\n",
    "\n",
    "df_rna = pd.read_csv(rna_path, sep='\\t')\n",
    "\n",
    "# setting gene ID as index\n",
    "if 'gene_id' in df_rna.columns:\n",
    "    df_rna = df_rna.set_index('gene_id')\n",
    "\n",
    "if 'tracking_id' in df_rna.columns:\n",
    "    df_rna = df_rna.drop(columns=['tracking_id'])\n",
    "\n",
    "# RNA Header ex: \"123_456_0\"  \n",
    "# Metadata ID ex: \"123_456\"\n",
    "valid_cols = []\n",
    "renamed_cols = {}\n",
    "\n",
    "for col in df_rna.columns:\n",
    "    cleaned_id = col.rsplit('_', 1)[0]\n",
    "    \n",
    "    if cleaned_id in valid_specimens:\n",
    "        valid_cols.append(col)\n",
    "        renamed_cols[col] = cleaned_id\n",
    "\n",
    "print(f\"Matched {len(valid_cols)} columns from RNA file to metadata\")\n",
    "\n",
    "df_rna = df_rna[valid_cols]\n",
    "df_rna = df_rna.rename(columns=renamed_cols)\n",
    "df_rna = df_rna.T\n",
    "\n",
    "print(f\"Final Cleaned RNA shape: {df_rna.shape}\")\n",
    "display(df_rna.head())\n",
    "\n",
    "# Should see:\n",
    "#Found 3196 valid specimenIDs in metadata\n",
    "#Matched 640 columns from RNA file to metadata\n",
    "#Final Cleaned RNA shape: (640, 55889)\n",
    "\n",
    "# we perform the same steps for Methylation\n",
    "\n",
    "meth_path = downloaded_files[\"ROSMAP_arrayMethylation_imputed.tsv.gz\"]\n",
    "meth_meta_path = downloaded_files[\"ROSMAP_assay_methylationArray_metadata.csv\"]\n",
    "\n",
    "df_meta_meth = pd.read_csv(meth_meta_path)\n",
    "valid_meth_ids = set(df_meta_meth['specimenID'].astype(str))\n",
    "\n",
    "print(f\"Found {len(valid_meth_ids)} valid specimenIDs in Methylation metadata\")\n",
    "\n",
    "df_meth = pd.read_csv(meth_path, sep='\\t', low_memory=False)\n",
    "df_meth = df_meth.set_index(df_meth.columns[0])\n",
    "\n",
    "print(f\"Raw Meth shape: {df_meth.shape}\")\n",
    "\n",
    "valid_cols = []\n",
    "for c in df_meth.columns:\n",
    "    if c in valid_meth_ids:\n",
    "        valid_cols.append(c)\n",
    "\n",
    "print(f\"Matched {len(valid_cols)} columns from Meth file to metadata\")\n",
    "\n",
    "df_meth = df_meth[valid_cols]\n",
    "df_meth = df_meth.T\n",
    "\n",
    "print(f\"Final Cleaned Meth shape: {df_meth.shape}\")\n",
    "display(df_meth.head())\n",
    "\n",
    "#Found 748 valid specimenIDs in Methylation metadata\n",
    "#Raw Meth shape: (420132, 740)\n",
    "#Matched 740 columns from Meth file to metadata\n",
    "#Final Cleaned Meth shape: (740, 420132)\n",
    "\n",
    "# Lastly we do the same for mirna\n",
    "\n",
    "mirna_path = downloaded_files[\"ROSMAP_arraymiRNA.gct\"]\n",
    "mirna_meta_path = downloaded_files[\"ROSMAP_assay_miRNAarray_nanostring_metadata.csv\"]\n",
    "\n",
    "df_meta = pd.read_csv(mirna_meta_path)\n",
    "id_map = dict(zip(df_meta['mirna_id'], df_meta['specimenID']))\n",
    "\n",
    "with open(mirna_path, 'r') as f:\n",
    "    first_line = f.readline()\n",
    "\n",
    "df_mirna = pd.read_csv(mirna_path, sep='\\t', skiprows=2)\n",
    "\n",
    "print(f\"Columns found: {list(df_mirna.columns[:5])}\")\n",
    "\n",
    "if 'Name' in df_mirna.columns:\n",
    "    df_mirna = df_mirna.set_index('Name')\n",
    "elif 'ProbeID' in df_mirna.columns:\n",
    "    df_mirna = df_mirna.set_index('ProbeID')\n",
    "else:\n",
    "    print(f\"not match found\")\n",
    "\n",
    "if 'Description' in df_mirna.columns:\n",
    "    df_mirna = df_mirna.drop(columns=['Description'])\n",
    "\n",
    "valid_cols = [c for c in df_mirna.columns if c in id_map]\n",
    "print(f\"Matched {len(valid_cols)} sample columns\")\n",
    "\n",
    "df_mirna = df_mirna[valid_cols]\n",
    "df_mirna = df_mirna.rename(columns=id_map)\n",
    "df_mirna = df_mirna.T\n",
    "\n",
    "print(f\"Final Cleaned miRNA shape: {df_mirna.shape}\")\n",
    "display(df_mirna.head())\n",
    "\n",
    "#Should see:\n",
    "#Matched 525 sample columns\n",
    "#Final Cleaned miRNA shape: (525, 309)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5e7d5154",
   "metadata": {},
   "outputs": [],
   "source": [
    "common_ids = df_rna.index.intersection(df_mirna.index)\n",
    "\n",
    "print(f\"RNA samples: {df_rna.shape[0]}\")\n",
    "print(f\"miRNA samples: {df_mirna.shape[0]}\")\n",
    "print(f\"Overlapping samples: {len(common_ids)}\")\n",
    "\n",
    "df_rna_aligned = df_rna.loc[common_ids]\n",
    "df_mirna_aligned = df_mirna.loc[common_ids]\n",
    "\n",
    "df_rna_aligned = df_rna_aligned.sort_index()\n",
    "df_mirna_aligned = df_mirna_aligned.sort_index()\n",
    "\n",
    "print(f\"Aligned RNA Shape: {df_rna_aligned.shape}\")\n",
    "print(f\"Aligned miRNA Shape: {df_mirna_aligned.shape}\")\n",
    "\n",
    "# verify they are identical\n",
    "if df_rna_aligned.index.equals(df_mirna_aligned.index):\n",
    "    print(\"Success: Indices are perfectly aligned\")\n",
    "else:\n",
    "    print(\"Warning: Indices do not match perfectly\")\n",
    "\n",
    "#RNA samples: 640\n",
    "#miRNA samples: 525\n",
    "#Overlapping samples: 525\n",
    "\n",
    "#Aligned RNA Shape: (527, 55889)\n",
    "#Aligned miRNA Shape: (525, 309)\n",
    "#Warning: Indices do not match perfectly\n",
    "\n",
    "dupes = df_rna_aligned.index[df_rna_aligned.index.duplicated()]\n",
    "print(f\"IDs with multiple entries in RNA: {dupes.tolist()}\")\n",
    "\n",
    "df_rna_aligned = df_rna_aligned.groupby(level=0).mean()\n",
    "\n",
    "print(f\"\\nAligned RNA Shape: {df_rna_aligned.shape}\")\n",
    "print(f\"Aligned miRNA Shape: {df_mirna_aligned.shape}\")\n",
    "\n",
    "if df_rna_aligned.shape[0] == df_mirna_aligned.shape[0]:\n",
    "    print(\"Perfect match 525 vs 525\")\n",
    "\n",
    "#IDs with multiple entries in RNA: ['xxxxx', 'xxxxxx']\n",
    "\n",
    "#Aligned RNA Shape: (525, 55889)\n",
    "#Aligned miRNA Shape: (525, 309)\n",
    "#Perfect match 525 vs 525"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6b146216",
   "metadata": {},
   "outputs": [],
   "source": [
    "df_bio = pd.read_csv(downloaded_files[\"ROSMAP_biospecimen_metadata.csv\"])\n",
    "\n",
    "# load metadata\n",
    "df_meta_rna = pd.read_csv(downloaded_files[\"ROSMAP_assay_RNAseq_metadata.csv\"])\n",
    "df_meta_meth = pd.read_csv(downloaded_files[\"ROSMAP_assay_methylationArray_metadata.csv\"])\n",
    "df_meta_mir = pd.read_csv(downloaded_files[\"ROSMAP_assay_miRNAarray_nanostring_metadata.csv\"])\n",
    "\n",
    "# keep just specimenID + individualID from biospecimen\n",
    "bio_key = df_bio[[\"specimenID\", \"individualID\"]].drop_duplicates()\n",
    "\n",
    "# attach ID to each df\n",
    "meta_rna_with_ind = df_meta_rna.merge(bio_key, on=\"specimenID\", how=\"left\")\n",
    "meta_meth_with_ind = df_meta_meth.merge(bio_key, on=\"specimenID\", how=\"left\")\n",
    "meta_mir_with_ind = df_meta_mir.merge(bio_key, on=\"specimenID\", how=\"left\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5ac82d3b",
   "metadata": {},
   "outputs": [],
   "source": [
    "rna_id_to_ind = (meta_rna_with_ind.drop_duplicates(\"specimenID\").set_index(\"specimenID\")[\"individualID\"])\n",
    "meth_id_to_ind = (meta_meth_with_ind.drop_duplicates(\"specimenID\").set_index(\"specimenID\")[\"individualID\"])\n",
    "mir_id_to_ind = (meta_mir_with_ind.drop_duplicates(\"specimenID\").set_index(\"specimenID\")[\"individualID\"])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7e7021ec",
   "metadata": {},
   "outputs": [],
   "source": [
    "# RNA\n",
    "rna_ind = df_rna_aligned.copy()\n",
    "rna_ind[\"individualID\"] = rna_id_to_ind.reindex(rna_ind.index).values\n",
    "rna_ind = rna_ind.dropna(subset=[\"individualID\"]).set_index(\"individualID\")\n",
    "\n",
    "# miRNA\n",
    "mirna_ind = df_mirna_aligned.copy()\n",
    "mirna_ind[\"individualID\"] = mir_id_to_ind.reindex(mirna_ind.index).values\n",
    "mirna_ind = mirna_ind.dropna(subset=[\"individualID\"]).set_index(\"individualID\")\n",
    "\n",
    "# Methylation: df_meth.index is methylation specimenID\n",
    "meth_ind = df_meth.copy()\n",
    "meth_ind[\"individualID\"] = meth_id_to_ind.reindex(meth_ind.index).values\n",
    "meth_ind = meth_ind.dropna(subset=[\"individualID\"]).set_index(\"individualID\")\n",
    "\n",
    "common_ids = sorted(\n",
    "    set(rna_ind.index) &\n",
    "    set(mirna_ind.index) &\n",
    "    set(meth_ind.index)\n",
    ")\n",
    "\n",
    "rna_final = rna_ind.loc[common_ids]\n",
    "mirna_final = mirna_ind.loc[common_ids]\n",
    "meth_final  = meth_ind.loc[common_ids]\n",
    "\n",
    "print(f\"RNA shape: {rna_final.shape}\")\n",
    "print(f\"miRNA shape: {mirna_final.shape}\")\n",
    "print(f\"Meth shape: {meth_final.shape}\")\n",
    "\n",
    "# should see\n",
    "#RNA shape: (521, 55889)\n",
    "#miRNA shape: (521, 309)\n",
    "#Meth shape: (521, 420132)\n",
    "\n",
    "display(meth_final.head(2))\n",
    "display(rna_final.head(2))\n",
    "display(meth_final.head(2))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "17738c00",
   "metadata": {},
   "outputs": [],
   "source": [
    "df_clin = pd.read_csv(downloaded_files[\"ROSMAP_clinical.csv\"])\n",
    "df_clin_ind = df_clin.set_index(\"individualID\").loc[common_ids]\n",
    "df_clin_ind.shape\n",
    "\n",
    "# this should show a DF of shape (521,17)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8f7f384b",
   "metadata": {},
   "outputs": [],
   "source": [
    "df_clin = pd.read_csv(downloaded_files[\"ROSMAP_clinical.csv\"])\n",
    "df_clin_ind = df_clin.set_index(\"individualID\")\n",
    "\n",
    "common_ids_all = sorted(set(common_ids) & set(df_clin_ind.index))\n",
    "\n",
    "rna_final = rna_final.loc[common_ids_all]\n",
    "mirna_final = mirna_final.loc[common_ids_all]\n",
    "meth_final = meth_final.loc[common_ids_all]\n",
    "clin_final = df_clin_ind.loc[common_ids_all]\n",
    "\n",
    "print(f\"RNA shape: {rna_final.shape}\")\n",
    "print(f\"miRNA shape: {mirna_final.shape}\")\n",
    "print(f\"Meth shape: {meth_final.shape}\")\n",
    "print(f\"Clinical shape: {clin_final.shape}\")\n",
    "\n",
    "# should see:\n",
    "#RNA shape: (521, 55889)\n",
    "#miRNA shape: (521, 309)\n",
    "#Meth shape: (521, 420132)\n",
    "#Clinical shape: (521, 17)\n",
    "\n",
    "display(clin_final.head())\n",
    "clin_final[\"cogdx\"].value_counts()\n",
    "\n",
    "# should see:\n",
    "# cogdx\n",
    "# 4.0    182\n",
    "# 1.0    168\n",
    "# 2.0    125\n",
    "# 5.0     25\n",
    "# 6.0     12\n",
    "# 3.0      9\n",
    "# Name: count, dtype: int64"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5e75e1b3",
   "metadata": {},
   "outputs": [],
   "source": [
    "from bioneuralnet import utils\n",
    "\n",
    "# This is a computationally expensive step: DO IT ONCE, then save result.\n",
    "print(\"\\nConverting meth beta-values to M-values\")\n",
    "X_meth_m = utils.m_transform(meth_final, eps=1e-6)\n",
    "\n",
    "print(\"Selecting top 100,000 by variance\")\n",
    "X_meth_100k = utils.variance_threshold(X_meth_m, k=100000)\n",
    "\n",
    "print(f\"X_meth_100k shape: {X_meth_100k.shape}\")\n",
    "\n",
    "# free big stuff if memory is limited\n",
    "del meth_final, X_meth_m"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "33f4905e",
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "\n",
    "labeled_patients = clin_final.index[clin_final[\"cogdx\"].notna()]\n",
    "print(\"Patients with non-missing cogdx:\", len(labeled_patients))\n",
    "\n",
    "common_labeled_patients = (\n",
    "    set(labeled_patients)\n",
    "    & set(X_meth_100k.index)\n",
    "    & set(rna_final.index)\n",
    "    & set(mirna_final.index)\n",
    ")\n",
    "common_labeled_patients = sorted(common_labeled_patients)\n",
    "print(\"Final labeled patients used for modeling:\", len(common_labeled_patients))\n",
    "\n",
    "X_meth = X_meth_100k.loc[common_labeled_patients]\n",
    "X_rna = rna_final.loc[common_labeled_patients]\n",
    "X_mirna = mirna_final.loc[common_labeled_patients]\n",
    "clinical_aligned = clin_final.loc[common_labeled_patients]\n",
    "\n",
    "cogdx_series = clinical_aligned[\"cogdx\"]\n",
    "\n",
    "def cogdx_to_3class(x):\n",
    "    if x == 1.0:\n",
    "        return 0\n",
    "    elif x in (2.0, 3.0):\n",
    "        return 1\n",
    "    elif x in (4.0, 5.0, 6.0):\n",
    "        return 2\n",
    "    else:\n",
    "        return np.nan\n",
    "    \n",
    "target = cogdx_series.map(cogdx_to_3class)\n",
    "valid_mask = target.notna()\n",
    "\n",
    "if not valid_mask.all():\n",
    "    print(f\"Dropping patients with invalid cogdx codes: {(~valid_mask).sum()}\")\n",
    "    final_patients = target.index[valid_mask]\n",
    "    X_meth = X_meth.loc[final_patients]\n",
    "    X_rna = X_rna.loc[final_patients]\n",
    "    X_mirna = X_mirna.loc[final_patients]\n",
    "    clinical_aligned = clinical_aligned.loc[final_patients]\n",
    "    target = target.loc[final_patients]\n",
    "else:\n",
    "    final_patients = list(target.index)\n",
    "\n",
    "target = target.astype(\"Int64\")\n",
    "y_df = pd.DataFrame(target.values, index=final_patients, columns=[\"target\"])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7757d920",
   "metadata": {},
   "outputs": [],
   "source": [
    "# check index is aligned\n",
    "assert X_meth.index.equals(X_rna.index)\n",
    "assert X_meth.index.equals(X_mirna.index)\n",
    "assert X_meth.index.equals(clinical_aligned.index)\n",
    "assert X_meth.index.equals(y_df.index)\n",
    "\n",
    "# same NAME for index in all dfs\n",
    "idx_name = X_meth.index.name or \"individualID\"\n",
    "X_meth.index.name = idx_name\n",
    "X_rna.index.name = idx_name\n",
    "X_mirna.index.name = idx_name\n",
    "clinical_aligned.index.name = idx_name\n",
    "y_df.index.name = idx_name\n",
    "\n",
    "print(\"\\nFinal shapes (aligned on same patients + index name):\")\n",
    "print(f\"X_meth: {X_meth.shape}\")\n",
    "print(f\"X_rna: {X_rna.shape}\")\n",
    "print(f\"X_mirna: {X_mirna.shape}\")\n",
    "print(f\"clinical:{clinical_aligned.shape}\")\n",
    "print(f\"y_df: {y_df.shape}\")\n",
    "\n",
    "print(\"\\nTarget value counts (0=NCI, 1=MCI, 2=Dementia):\")\n",
    "print(y_df['target'].value_counts())"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0c9a37af",
   "metadata": {},
   "source": [
    "## Verify the following output before proceeding:\n",
    "\n",
    "Final shapes (aligned on same patients + index name):\n",
    "\n",
    "- X_meth: (521, 100000)\n",
    "- X_rna: (521, 55889)\n",
    "- X_mirna: (521, 309)\n",
    "- clinical: (521, 17)\n",
    "- y_df: (521, 1)\n",
    "\n",
    "Target value counts (0=NCI, 1=MCI, 2=Dementia):\n",
    "\n",
    "- 2: 219\n",
    "- 0: 168\n",
    "- 1: 134\n",
    "\n",
    "Name: count, dtype: Int64"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7ce699fc",
   "metadata": {},
   "outputs": [],
   "source": [
    "# optional save in case we restart kernel\n",
    "# X_meth.to_parquet(\"X_meth_final.parquet\")\n",
    "# X_rna.to_parquet(\"X_rna_final.parquet\")\n",
    "# X_mirna.to_parquet(\"X_mirna_final.parquet\")\n",
    "# clinical_aligned.to_parquet(\"X_clinical_final.parquet\")\n",
    "# y_df.to_parquet(\"y_df_final.parquet\")\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e7f335a3",
   "metadata": {},
   "outputs": [],
   "source": [
    "# if saved then data can be loaded this way\n",
    "\n",
    "# X_meth = pd.read_parquet(\"X_meth_final.parquet\")\n",
    "# X_rna = pd.read_parquet(\"X_rna_final.parquet\")\n",
    "# X_mirna = pd.read_parquet(\"X_mirna_final.parquet\")\n",
    "# clinical = pd.read_parquet(\"X_clinical_final.parquet\")\n",
    "# y_df = pd.read_parquet(\"y_df_final.parquet\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7179717f",
   "metadata": {},
   "outputs": [],
   "source": [
    "from bioneuralnet import utils\n",
    "\n",
    "# Use the 3-class target as a Series\n",
    "Y_series = y_df[\"target\"].astype(int)\n",
    "\n",
    "print(\"\\nMethylation: ANOVA F-test\")\n",
    "meth_af = utils.top_anova_f_features(\n",
    "    X_meth,\n",
    "    Y_series,\n",
    "    max_features=8000\n",
    ")\n",
    "\n",
    "print(\"Methylation: RandomForest\")\n",
    "meth_rf = utils.importance_rf(\n",
    "    X_meth,\n",
    "    Y_series,\n",
    "    top_k=8000\n",
    ")\n",
    "\n",
    "meth_anova_set = set(meth_af.columns)\n",
    "meth_rf_set = set(meth_rf.columns)\n",
    "meth_inter_cols = sorted(meth_anova_set & meth_rf_set)\n",
    "\n",
    "print(f\"Methylation intersection size: {len(meth_inter_cols)}\")\n",
    "\n",
    "# methylation final shape after feature selection\n",
    "X_meth_fs = X_meth[meth_inter_cols]\n",
    "print(f\"X_meth_fs shape: {X_meth_fs.shape}\")\n",
    "\n",
    "print(\"\\nRNA: ANOVA F-test\")\n",
    "rna_af = utils.top_anova_f_features(\n",
    "    X_rna,\n",
    "    Y_series,\n",
    "    max_features=8000\n",
    ")\n",
    "\n",
    "print(\"RNA: RandomForest\")\n",
    "rna_rf = utils.importance_rf(\n",
    "    X_rna,\n",
    "    Y_series,\n",
    "    top_k=8000\n",
    ")\n",
    "\n",
    "rna_anova_set = set(rna_af.columns)\n",
    "rna_rf_set = set(rna_rf.columns)\n",
    "rna_inter_cols = sorted(rna_anova_set & rna_rf_set)\n",
    "\n",
    "print(f\"RNA intersection size: {len(rna_inter_cols)}\")\n",
    "\n",
    "X_rna_fs = X_rna[rna_inter_cols]\n",
    "print(f\"X_rna_fs shape: {X_rna_fs.shape}\")\n",
    "\n",
    "# At this stage we have the following shapes\n",
    "# X_meth: (521, 772)\n",
    "# X_rna : (521, 2054)\n",
    "# X_mirna: (521, 309)\n",
    "# clinical: (521, 17)\n",
    "# Y_labels: (521, 1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "eb9ef4cc",
   "metadata": {},
   "outputs": [],
   "source": [
    "from bioneuralnet.utils import preprocess_clinical\n",
    "\n",
    "clinical_for_model = preprocess_clinical(\n",
    "    clinical_aligned, \n",
    "    # internal variance feature selection\n",
    "    top_k=8, \n",
    "    scale=False,\n",
    "    ignore_columns=[\"projid\", \"dcfdx_lv\", \"cogdx\", \"braaksc\", \"ceradsc\", \"Study\", \"age_first_ad_dx\"],\n",
    "    nan_threshold=0.5)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ed9d1801",
   "metadata": {},
   "outputs": [],
   "source": [
    "import re\n",
    "\n",
    "def clean_bn_name(col: str) -> str:\n",
    "    col = str(col)\n",
    "    node_clean = re.sub(r\"[^0-9a-zA-Z_]\", \".\", col)\n",
    "    if not node_clean[0].isalpha():\n",
    "        node_clean = \"X\" + node_clean\n",
    "    return node_clean\n",
    "\n",
    "X_meth_clean = X_meth_fs.copy()\n",
    "meth_cols = []\n",
    "for c in X_meth_clean.columns:\n",
    "    meth_cols.append(clean_bn_name(c))\n",
    "X_meth_clean.columns = meth_cols\n",
    "\n",
    "# 2. Process X_rna_clean\n",
    "X_rna_clean = X_rna_fs.copy()\n",
    "rna_cols = []\n",
    "for c in X_rna_clean.columns:\n",
    "    rna_cols.append(clean_bn_name(c))\n",
    "X_rna_clean.columns = rna_cols\n",
    "\n",
    "# 3. Process X_mirna_clean\n",
    "X_mirna_clean = X_mirna.copy()\n",
    "mirna_cols = []\n",
    "for c in X_mirna_clean.columns:\n",
    "    mirna_cols.append(clean_bn_name(c))\n",
    "X_mirna_clean.columns = mirna_cols\n",
    "\n",
    "omics_list = [X_meth_clean, X_rna_clean, X_mirna_clean]\n",
    "\n",
    "omics_for_graph = pd.concat(omics_list, axis=1)\n",
    "print(omics_for_graph.shape)\n",
    "\n",
    "# Save everything in case kernel crashes\n",
    "omics_for_graph.to_parquet(\"omics_for_graph.parquet\")\n",
    "X_meth_clean.to_parquet(\"X_meth_clean.parquet\")\n",
    "X_rna_clean.to_parquet(\"X_rna_clean.parquet\")\n",
    "X_mirna_clean.to_parquet(\"X_mirna_clean.parquet\")\n",
    "\n",
    "clinical_for_model.to_parquet(\"clinical_for_model.parquet\")\n",
    "y_df.to_parquet(\"y_df.parquet\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d32d62ef",
   "metadata": {},
   "outputs": [],
   "source": [
    "print(f\"Selecting top 300 features from Methylation\")\n",
    "X_meth_clean3 = utils.variance_threshold(X_meth_clean, k=300)\n",
    "\n",
    "print(f\"Selecting top 300 features from RNA\")\n",
    "X_rna_clean3 = utils.variance_threshold(X_rna_clean, k=300)\n",
    "\n",
    "print(f\"Selecting top 300 features from miRNA\")\n",
    "X_mirna_clean3 = utils.variance_threshold(X_mirna_clean, k=300)\n",
    "\n",
    "omics_for_graph_selected = pd.concat([X_meth_clean3, X_rna_clean3, X_mirna_clean3], axis=1)\n",
    "new_total_features = X_meth_clean3.shape[1] + X_rna_clean3.shape[1] + X_mirna_clean3.shape[1]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7238da83",
   "metadata": {},
   "outputs": [],
   "source": [
    "from bioneuralnet.clustering import HybridLouvain\n",
    "from bioneuralnet.network import similarity_network\n",
    "\n",
    "similarity_10 = similarity_network(omics_for_graph_selected, k=10)\n",
    "\n",
    "SEED = 1883\n",
    "utils.set_seed(SEED)\n",
    "\n",
    "# NCI vs rest\n",
    "y_binary = (y_df[\"target\"] == 0).astype(int)\n",
    "\n",
    "hl = HybridLouvain(\n",
    "    G=similarity_10,\n",
    "    B=omics_for_graph,\n",
    "    Y=y_binary,\n",
    "    k_L=0.2,\n",
    "    k_P=0.8,\n",
    "    max_iter=5,\n",
    "    seed=SEED,\n",
    ")\n",
    "\n",
    "subnetworks = hl.run(as_dfs=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "404fcd0a",
   "metadata": {},
   "outputs": [],
   "source": [
    "from bioneuralnet.metrics import plot_network\n",
    "\n",
    "nci_submodule = subnetworks[3]\n",
    "print(f\"NCI vs Rest submodule: {nci_submodule.shape[1]} features\")\n",
    "\n",
    "nci_mapping = plot_network(nci_submodule, weight_threshold=.45, show_labels=True, show_edge_weights=False)\n",
    "print(nci_mapping)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "32bde8d8",
   "metadata": {},
   "outputs": [],
   "source": [
    "from bioneuralnet.network import NetworkAnalyzer\n",
    "\n",
    "nci_submodule_50 = subnetworks[1]\n",
    "nci_submodule_25 = subnetworks[2]\n",
    "\n",
    "nci50_net = NetworkAnalyzer(nci_submodule_50)\n",
    "nci50_net.basic_statistics()\n",
    "\n",
    "nci25_net = NetworkAnalyzer(nci_submodule_25)\n",
    "nci25_net.basic_statistics()\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ed29ee7d",
   "metadata": {},
   "outputs": [],
   "source": [
    "from bioneuralnet import DPMON\n",
    "\n",
    "y_binary = (y_df[\"target\"] == 0).astype(int)\n",
    "\n",
    "# no tunning: using default parameters. ONly goal is to generate emebddings\n",
    "dpmon_params_base = {\n",
    "    \"adjacency_matrix\": similarity_10,\n",
    "    \"omics_list\": omics_for_graph_selected,\n",
    "    \"phenotype_data\": y_binary,\n",
    "    \"phenotype_col\": \"target\",\n",
    "    \"clinical_data\": clinical_for_model,\n",
    "\n",
    "}\n",
    "\n",
    "dpmon_obj = DPMON(**dpmon_params_base)\n",
    "_, _, embeddings = dpmon_obj.run()\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0d2ea72b",
   "metadata": {},
   "outputs": [],
   "source": [
    "# .detach().cpu().numpy() to get the raw numbers out of pytorch\n",
    "GAT_embeddings_df = pd.DataFrame(\n",
    "    embeddings.detach().cpu().numpy(),\n",
    "    index=similarity_10.columns\n",
    ")\n",
    "\n",
    "display(GAT_embeddings_df.iloc[:5])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c94debfe",
   "metadata": {},
   "outputs": [],
   "source": [
    "from bioneuralnet.metrics import plot_embeddings\n",
    "\n",
    "node_labels = []\n",
    "GAT_embeddings_array = GAT_embeddings_df.values\n",
    "feature_names = GAT_embeddings_df.index\n",
    "\n",
    "dna_feats = set(X_meth_clean3.columns)\n",
    "rna_feats = set(X_rna_clean3.columns)\n",
    "mirna_feats = set(X_mirna_clean3.columns)\n",
    "\n",
    "\n",
    "for feat in feature_names:\n",
    "    if feat in dna_feats:\n",
    "        # 1 for DNA Methylation\n",
    "        node_labels.append(1)\n",
    "    elif feat in rna_feats:\n",
    "        # 2 for RNA \n",
    "        node_labels.append(2)\n",
    "    elif feat in mirna_feats:\n",
    "        # 3 for miRNA\n",
    "        node_labels.append(3)\n",
    "    else:\n",
    "        # 0 for everything else, just in case\n",
    "        node_labels.append(0)\n",
    "\n",
    "# our plotting function needs an array to squeeze the embedings into 2Dspace\n",
    "node_labels = np.array(node_labels)\n",
    "\n",
    "plot_embeddings(GAT_embeddings_array, node_labels, legend_labels=[\"Background\", \"NCI_Meth\", \"NCI_RNA\", \"NCI_miRNA\"])"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8def20c5",
   "metadata": {},
   "source": [
    "## References:\n",
    "\n",
    "AD Knowledge Portal. Available from: [AD Knowledge Portal](https://adknowledgeportal.org/). Data generated from postmortem brain tissue provided by the Religious Orders Study and Rush Memory and Aging Project (ROSMAP) cohort at Rush Alzheimer’s Disease Center, Rush University Medical Center, Chicago. Supported by Cure Alzheimer’s Fund and NIH grants AG058002, AG062377, NS110453, NS115064, AG062335, AG074003, NS127187, MH119509, HG008155, RF1AG062377, RF1AG054321, R01AG054012, and GM087237."
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": ".enviroment",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.12.3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}