04_calculate_SAGE_values_VAE.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from models import *\n",
    "import sys\n",
    "\n",
    "import numpy as np\n",
    "import pandas as pd\n",
    "\n",
    "import sage\n",
    "\n",
    "\n",
    "############################################\n",
    "# Define helper functions\n",
    "############################################\n",
    "\n",
    "def save_sages(sage_vals, path):\n",
    "    res = pd.DataFrame([aml_data.columns.to_list(), sage_vals.values, sage_vals.std]).T\n",
    "    res.columns = [\"metabolite_id\", \"sage_value\", \"sage_value_sd\"]\n",
    "    res.to_csv(path, index=False)\n",
    "    \n",
    "def save_pw_sages(sage_vals, pw_groupnames, path):\n",
    "    res = pd.DataFrame([pw_groupnames, sage_vals.values, sage_vals.std]).T\n",
    "    res.columns = [\"pathway_name\", \"sage_value\", \"sage_value_sd\"]\n",
    "    res.to_csv(path, index=False)\n",
    "\n",
    "def get_sage_pws(model, ref_data, test_data, dim, pw_groups, pw_groupnames, pw_type):\n",
    "\n",
    "    def get_dim_vals(dat):\n",
    "        return model.encode_mu(dat)[:,dim:(dim+1)]\n",
    "    \n",
    "    # calculate base values\n",
    "    dim_output = get_dim_vals(test_data)\n",
    "    \n",
    "    # Setup and calculate\n",
    "    # NOTE: any callable function that returns a prediction is allowed in PermutationSampler\n",
    "    imputer_pw = sage.GroupedMarginalImputer(ref_data[:10], pw_groups)\n",
    "    sampler_pw = sage.PermutationSampler(get_dim_vals, imputer_pw, 'mse')\n",
    "    sage_values_pw = sampler_pw(test_data, Y=dim_output, batch_size = 512)\n",
    "    \n",
    "    # rename dimension 0 to dimension 18\n",
    "    dim_idx = 18 if dim == 0 else dim\n",
    "    \n",
    "    if pw_type == \"super\": \n",
    "        save_pw_sages(sage_values_pw, pw_groupnames, f\"results/sage_values/VAE/superpw_dim_{dim_idx}.csv\")\n",
    "    elif pw_type == \"sub\":\n",
    "        save_pw_sages(sage_values_pw, pw_groupnames, f\"results/sage_values/VAE/subpw_dim_{dim_idx}.csv\")\n",
    "        \n",
    "    return sage_values_pw\n",
    "\n",
    "\n",
    "def get_sage_mets(model, ref_data, test_data, dim):\n",
    "\n",
    "    def get_dim_vals(dat):\n",
    "        return model.encode_mu(dat)[:,dim:(dim+1)]\n",
    "    \n",
    "    dim_output = get_dim_vals(test_data)\n",
    "    \n",
    "    # Setup and calculate\n",
    "    # NOTE: any callable function that returns a prediction is allowed in PermutationSampler\n",
    "    imputer = sage.MarginalImputer(ref_data[:10])\n",
    "    sampler = sage.PermutationSampler(get_dim_vals, imputer, 'mse')\n",
    "    sage_values = sampler(test_data, Y=dim_output, batch_size = 10)\n",
    "    \n",
    "    # rename dimension 0 to dimension 18\n",
    "    dim_idx = 18 if dim == 0 else dim\n",
    "    \n",
    "    save_sages(sage_values, f\"results/sage_values/VAE/met_dim_{dim_idx}.csv\")\n",
    "    \n",
    "    return sage_values\n",
    "\n",
    "\n",
    "def get_groups(pw_name):\n",
    "    # Feature groups\n",
    "    feature_groups = met_annos.groupby(pw_name)['COMP_IDstr'].apply(list).to_dict()\n",
    "\n",
    "    group_names = [group for group in feature_groups]\n",
    "    for col in feature_names:\n",
    "        if np.all([col not in group[1] for group in feature_groups.items()]):\n",
    "            group_names.append(col)\n",
    "\n",
    "    # Group indices\n",
    "    groups = []\n",
    "    for _, group in feature_groups.items():\n",
    "        ind_list = []\n",
    "        for feature in group:\n",
    "            ind_list.append(data_cols.index(feature))\n",
    "        groups.append(ind_list)\n",
    "        \n",
    "    return {'feature_groups':feature_groups, 'group_names':group_names, 'groups':groups}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "############################################\n",
    "# Initialize variables and instantiate objects\n",
    "############################################\n",
    "\n",
    "twins_path  = 'data/TwinsUK.xls'\n",
    "aml_path    = 'data/AML.xls'\n",
    "\n",
    "aml_data = pd.read_excel(aml_path, sheet_name='Metabolite Data')\n",
    "aml_anno = pd.read_excel(aml_path, sheet_name='Sample Annotations')\n",
    "\n",
    "twins_train_df = pd.read_excel(twins_path, sheet_name='Training Set')\n",
    "twins_test_df  = pd.read_excel(twins_path, sheet_name='Testing Set')\n",
    "\n",
    "# these arrays are used for score calculations\n",
    "twins_train = twins_train_df.values\n",
    "twins_test  = twins_test_df.values\n",
    "\n",
    "\n",
    "model_path = 'models/'\n",
    "\n",
    "# Define paths for saved VAE model\n",
    "path_vae     = model_path + 'VAE_weights.h5'\n",
    "path_encoder = model_path + 'encoder_weights.h5'\n",
    "path_decoder = model_path + 'decoder_weights.h5'\n",
    "\n",
    "\n",
    "# Data & model configuration\n",
    "input_dim = twins_train_df.shape[1]\n",
    "intermediate_dim = 200\n",
    "latent_dim = 18\n",
    "\n",
    "kl_beta = 1e-2\n",
    "learning_rate = 1e-3\n",
    "\n",
    "# instantiate model\n",
    "mtmodel = mtVAE(input_dim,\n",
    "                intermediate_dim,\n",
    "                latent_dim,\n",
    "                kl_beta,\n",
    "                learning_rate)\n",
    "\n",
    "# load model\n",
    "mtmodel.vae.load_weights(path_vae)\n",
    "mtmodel.encoder.load_weights(path_encoder)\n",
    "mtmodel.decoder.load_weights(path_decoder)\n",
    "\n",
    "\n",
    "# Create groups to calculate grouped SAGE values\n",
    "data_cols = aml_data.columns.to_list()\n",
    "met_annos = pd.read_excel(aml_path, sheet_name='Metabolite Annotations')\n",
    "met_annos = met_annos.loc[met_annos['COMP_IDstr'].isin(data_cols)]\n",
    "feature_names = met_annos['COMP_IDstr'].to_list()\n",
    "biochemical_names = met_annos['BIOCHEMICAL'].to_list()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "############################################\n",
    "# MAIN PART OF SCRIPT\n",
    "############################################\n",
    "\n",
    "for latent_dim in range(18):\n",
    "    for sage_type in ['superpathway', 'subpathway', 'metabolite']:\n",
    "        if sage_type == \"superpathway\":\n",
    "            superpw        = get_groups('SUPER_PATHWAY')\n",
    "            sage_pw_values = get_sage_pws(mtmodel, twins_train, twins_test, latent_dim, superpw['groups'], superpw['group_names'], \"super\")\n",
    "\n",
    "        elif sage_type == \"subpathway\":\n",
    "            subpw          = get_groups('SUB_PATHWAY')\n",
    "            sage_pw_values = get_sage_pws(mtmodel, twins_train, twins_test, latent_dim, subpw['groups'], subpw['group_names'], \"sub\")\n",
    "\n",
    "        elif sage_type == \"metabolite\":\n",
    "            sage_met_values = get_sage_mets(mtmodel, twins_train, twins_test, latent_dim)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "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.6.12"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}