2.train-VAE/6.apply_nmf.ipynb

{
 "cells": [
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "fdc3df1b-9f13-464f-94b3-b088c804965d",
   "metadata": {},
   "source": [
    "## Apply NMF to gene dependencies\n",
    "\n",
    "Non-negative Matrix Factorization (NMF) is a dimensionality reduction technique that factors a non-negative matrix into two non-negative matrices. This is particularly useful when the data is inherently non-negative.\n",
    "\n",
    "We apply it to GeneEffect scores here and save latent representations (NMF components) for downstream comparative analyses (to compare with BetaVAE)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "e9edbfb3-e84e-4435-9900-4def2649cc6c",
   "metadata": {},
   "outputs": [],
   "source": [
    "import sys\n",
    "import pathlib\n",
    "import pandas as pd\n",
    "import plotnine as gg\n",
    "from sklearn.decomposition import NMF\n",
    "\n",
    "sys.path.insert(0, \"../utils/\")\n",
    "from data_loader import load_model_data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "9259847e-982d-4486-ab72-de50888aae36",
   "metadata": {},
   "outputs": [],
   "source": [
    "data_directory = pathlib.Path(\"../0.data-download/data\").resolve()\n",
    "dependency_file = pathlib.Path(f\"{data_directory}/CRISPRGeneEffect.parquet\").resolve()\n",
    "gene_dict_file = pathlib.Path(f\"{data_directory}/CRISPR_gene_dictionary.parquet\").resolve()\n",
    "\n",
    "\n",
    "output_dir = pathlib.Path(\"results\").resolve()\n",
    "nmf_output_file = pathlib.Path(f\"{output_dir}/nmf_latent.parquet.gz\").resolve()\n",
    "output_nmf_weights_file = pathlib.Path(f\"{output_dir}/NMF_weight_matrix_gsea.parquet\").resolve()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "4b736280-f8f3-46c6-a229-7f2d757f6f54",
   "metadata": {},
   "outputs": [],
   "source": [
    "nmf_components = 50"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "09c9c59c-5be6-4fca-96b3-c7f303a20b34",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "(1150, 18444)\n"
     ]
    }
   ],
   "source": [
    "# Load data\n",
    "dependency_df, gene_dict_df = load_model_data(dependency_file, gene_dict_file)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "dc5ba3e1",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Ensure all values are non-negative by shifting the data\n",
    "min_value = dependency_df.drop(columns=[\"ModelID\"]).min().min()\n",
    "if min_value < 0:\n",
    "   dependency_df_non_negative = dependency_df.drop(columns=[\"ModelID\"]) - min_value\n",
    "else:\n",
    "   dependency_df_non_negative = dependency_df.drop(columns=[\"ModelID\"])"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "fe75f876-9b6b-4444-bc0b-0c67135f955e",
   "metadata": {},
   "source": [
    "# Perform NMF"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "f7d590b4-23c9-4f92-820a-8fbe1fe3c7c9",
   "metadata": {},
   "outputs": [],
   "source": [
    "nmf = NMF(n_components=nmf_components, random_state=0)\n",
    "W = nmf.fit_transform(dependency_df_non_negative)\n",
    "H = nmf.components_"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "65756252",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Transform models into NMF space (W matrix)\n",
    "dependency_df_transformed = pd.DataFrame(W)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "12ee4081-4d2a-4266-96a4-334b985871b0",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "(1150, 51)\n"
     ]
    },
    {
     "data": {
      "text/html": [
       "<div>\n",
       "<style scoped>\n",
       "    .dataframe tbody tr th:only-of-type {\n",
       "        vertical-align: middle;\n",
       "    }\n",
       "\n",
       "    .dataframe tbody tr th {\n",
       "        vertical-align: top;\n",
       "    }\n",
       "\n",
       "    .dataframe thead th {\n",
       "        text-align: right;\n",
       "    }\n",
       "</style>\n",
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th></th>\n",
       "      <th>ModelID</th>\n",
       "      <th>NMF_0</th>\n",
       "      <th>NMF_1</th>\n",
       "      <th>NMF_2</th>\n",
       "      <th>NMF_3</th>\n",
       "      <th>NMF_4</th>\n",
       "      <th>NMF_5</th>\n",
       "      <th>NMF_6</th>\n",
       "      <th>NMF_7</th>\n",
       "      <th>NMF_8</th>\n",
       "      <th>...</th>\n",
       "      <th>NMF_40</th>\n",
       "      <th>NMF_41</th>\n",
       "      <th>NMF_42</th>\n",
       "      <th>NMF_43</th>\n",
       "      <th>NMF_44</th>\n",
       "      <th>NMF_45</th>\n",
       "      <th>NMF_46</th>\n",
       "      <th>NMF_47</th>\n",
       "      <th>NMF_48</th>\n",
       "      <th>NMF_49</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>ACH-000001</td>\n",
       "      <td>0.209800</td>\n",
       "      <td>0.000000</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.014303</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.000000</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.000000</td>\n",
       "      <td>0.000000</td>\n",
       "      <td>...</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.009268</td>\n",
       "      <td>0.004850</td>\n",
       "      <td>0.982635</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1</th>\n",
       "      <td>ACH-000004</td>\n",
       "      <td>0.215258</td>\n",
       "      <td>0.004779</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.000000</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.004609</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.002082</td>\n",
       "      <td>0.000000</td>\n",
       "      <td>...</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.001776</td>\n",
       "      <td>0.012277</td>\n",
       "      <td>0.973513</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>2</th>\n",
       "      <td>ACH-000005</td>\n",
       "      <td>0.215434</td>\n",
       "      <td>0.000937</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.000000</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.012313</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.000643</td>\n",
       "      <td>0.000318</td>\n",
       "      <td>...</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.002394</td>\n",
       "      <td>0.013144</td>\n",
       "      <td>0.974102</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "<p>3 rows × 51 columns</p>\n",
       "</div>"
      ],
      "text/plain": [
       "      ModelID     NMF_0     NMF_1  NMF_2     NMF_3  NMF_4     NMF_5  NMF_6  \\\n",
       "0  ACH-000001  0.209800  0.000000    0.0  0.014303    0.0  0.000000    0.0   \n",
       "1  ACH-000004  0.215258  0.004779    0.0  0.000000    0.0  0.004609    0.0   \n",
       "2  ACH-000005  0.215434  0.000937    0.0  0.000000    0.0  0.012313    0.0   \n",
       "\n",
       "      NMF_7     NMF_8  ...  NMF_40  NMF_41  NMF_42  NMF_43  NMF_44  NMF_45  \\\n",
       "0  0.000000  0.000000  ...     0.0     0.0     0.0     0.0     0.0     0.0   \n",
       "1  0.002082  0.000000  ...     0.0     0.0     0.0     0.0     0.0     0.0   \n",
       "2  0.000643  0.000318  ...     0.0     0.0     0.0     0.0     0.0     0.0   \n",
       "\n",
       "   NMF_46    NMF_47    NMF_48    NMF_49  \n",
       "0     0.0  0.009268  0.004850  0.982635  \n",
       "1     0.0  0.001776  0.012277  0.973513  \n",
       "2     0.0  0.002394  0.013144  0.974102  \n",
       "\n",
       "[3 rows x 51 columns]"
      ]
     },
     "execution_count": 8,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# Recode column space and add back model IDs\n",
    "dependency_df_transformed.columns = [f\"NMF_{x}\" for x in range(0, dependency_df_transformed.shape[1])]\n",
    "dependency_df_transformed = pd.concat([dependency_df.loc[:, \"ModelID\"], dependency_df_transformed], axis=\"columns\")\n",
    "\n",
    "\n",
    "dependency_df_transformed.to_parquet(nmf_output_file, index=False)\n",
    "\n",
    "\n",
    "print(dependency_df_transformed.shape)\n",
    "dependency_df_transformed.head(3)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "26ab257b",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "(2718, 51)\n"
     ]
    },
    {
     "data": {
      "text/html": [
       "<div>\n",
       "<style scoped>\n",
       "    .dataframe tbody tr th:only-of-type {\n",
       "        vertical-align: middle;\n",
       "    }\n",
       "\n",
       "    .dataframe tbody tr th {\n",
       "        vertical-align: top;\n",
       "    }\n",
       "\n",
       "    .dataframe thead th {\n",
       "        text-align: right;\n",
       "    }\n",
       "</style>\n",
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th></th>\n",
       "      <th>genes</th>\n",
       "      <th>NMF_0</th>\n",
       "      <th>NMF_1</th>\n",
       "      <th>NMF_2</th>\n",
       "      <th>NMF_3</th>\n",
       "      <th>NMF_4</th>\n",
       "      <th>NMF_5</th>\n",
       "      <th>NMF_6</th>\n",
       "      <th>NMF_7</th>\n",
       "      <th>NMF_8</th>\n",
       "      <th>...</th>\n",
       "      <th>NMF_40</th>\n",
       "      <th>NMF_41</th>\n",
       "      <th>NMF_42</th>\n",
       "      <th>NMF_43</th>\n",
       "      <th>NMF_44</th>\n",
       "      <th>NMF_45</th>\n",
       "      <th>NMF_46</th>\n",
       "      <th>NMF_47</th>\n",
       "      <th>NMF_48</th>\n",
       "      <th>NMF_49</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>SARS2</td>\n",
       "      <td>0.071005</td>\n",
       "      <td>0.000000</td>\n",
       "      <td>21.368498</td>\n",
       "      <td>11.065316</td>\n",
       "      <td>0.000000</td>\n",
       "      <td>7.578911</td>\n",
       "      <td>0.000000</td>\n",
       "      <td>29.019315</td>\n",
       "      <td>61.728789</td>\n",
       "      <td>...</td>\n",
       "      <td>0.000000</td>\n",
       "      <td>0.00000</td>\n",
       "      <td>9.783745</td>\n",
       "      <td>4.862253</td>\n",
       "      <td>1.096877</td>\n",
       "      <td>3.659648</td>\n",
       "      <td>3.391358</td>\n",
       "      <td>3.903741</td>\n",
       "      <td>3.062480</td>\n",
       "      <td>5.001846</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1</th>\n",
       "      <td>LARP4</td>\n",
       "      <td>0.322014</td>\n",
       "      <td>16.745962</td>\n",
       "      <td>7.208615</td>\n",
       "      <td>3.910373</td>\n",
       "      <td>16.466121</td>\n",
       "      <td>0.478806</td>\n",
       "      <td>13.895235</td>\n",
       "      <td>10.180814</td>\n",
       "      <td>31.927719</td>\n",
       "      <td>...</td>\n",
       "      <td>25.773102</td>\n",
       "      <td>3.64141</td>\n",
       "      <td>0.000000</td>\n",
       "      <td>0.901202</td>\n",
       "      <td>5.477956</td>\n",
       "      <td>3.124627</td>\n",
       "      <td>4.322935</td>\n",
       "      <td>3.423482</td>\n",
       "      <td>3.464876</td>\n",
       "      <td>5.399670</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>2</th>\n",
       "      <td>APH1A</td>\n",
       "      <td>4.372661</td>\n",
       "      <td>3.957601</td>\n",
       "      <td>0.000000</td>\n",
       "      <td>0.000000</td>\n",
       "      <td>0.000000</td>\n",
       "      <td>0.000000</td>\n",
       "      <td>0.000000</td>\n",
       "      <td>0.000000</td>\n",
       "      <td>0.000000</td>\n",
       "      <td>...</td>\n",
       "      <td>0.000000</td>\n",
       "      <td>0.00000</td>\n",
       "      <td>7.401898</td>\n",
       "      <td>4.173206</td>\n",
       "      <td>3.360864</td>\n",
       "      <td>2.882472</td>\n",
       "      <td>1.783110</td>\n",
       "      <td>3.483496</td>\n",
       "      <td>4.298280</td>\n",
       "      <td>4.546053</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "<p>3 rows × 51 columns</p>\n",
       "</div>"
      ],
      "text/plain": [
       "   genes     NMF_0      NMF_1      NMF_2      NMF_3      NMF_4     NMF_5  \\\n",
       "0  SARS2  0.071005   0.000000  21.368498  11.065316   0.000000  7.578911   \n",
       "1  LARP4  0.322014  16.745962   7.208615   3.910373  16.466121  0.478806   \n",
       "2  APH1A  4.372661   3.957601   0.000000   0.000000   0.000000  0.000000   \n",
       "\n",
       "       NMF_6      NMF_7      NMF_8  ...     NMF_40   NMF_41    NMF_42  \\\n",
       "0   0.000000  29.019315  61.728789  ...   0.000000  0.00000  9.783745   \n",
       "1  13.895235  10.180814  31.927719  ...  25.773102  3.64141  0.000000   \n",
       "2   0.000000   0.000000   0.000000  ...   0.000000  0.00000  7.401898   \n",
       "\n",
       "     NMF_43    NMF_44    NMF_45    NMF_46    NMF_47    NMF_48    NMF_49  \n",
       "0  4.862253  1.096877  3.659648  3.391358  3.903741  3.062480  5.001846  \n",
       "1  0.901202  5.477956  3.124627  4.322935  3.423482  3.464876  5.399670  \n",
       "2  4.173206  3.360864  2.882472  1.783110  3.483496  4.298280  4.546053  \n",
       "\n",
       "[3 rows x 51 columns]"
      ]
     },
     "execution_count": 9,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# Obtain weights (H matrix), which can be used in GSEA\n",
    "nmf_weights = pd.DataFrame(H, columns=dependency_df_non_negative.columns.tolist()).transpose()\n",
    "nmf_weights.columns = [f\"NMF_{x}\" for x in range(nmf_weights.shape[1])]\n",
    "\n",
    "\n",
    "nmf_weights = nmf_weights.reset_index().rename(columns={\"index\": \"genes\"})\n",
    "\n",
    "\n",
    "nmf_weights.to_parquet(output_nmf_weights_file, index=False)\n",
    "\n",
    "\n",
    "print(nmf_weights.shape)\n",
    "nmf_weights.head(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.12.4"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}