1_fit_pls.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Load libraries"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import pandas as pd\n",
    "import scanpy as sc\n",
    "import matplotlib.pyplot as plt\n",
    "import seaborn as sns\n",
    "\n",
    "from joblib import Parallel, delayed\n",
    "from numpy.random import default_rng\n",
    "\n",
    "from sklearn.cross_decomposition import PLSRegression\n",
    "from sklearn.model_selection import GridSearchCV, KFold, cross_val_score, cross_val_predict\n",
    "from sklearn.utils import shuffle\n",
    "from sklearn.preprocessing import StandardScaler\n",
    "from sklearn.decomposition import PCA\n",
    "from sklearn.pipeline import make_pipeline\n",
    "\n",
    "from scipy.stats import spearmanr, pearsonr\n",
    "\n",
    "import seaborn as sns\n",
    "\n",
    "from utils import loess_ci\n",
    "\n",
    "import gseapy as gp\n",
    "lib_names = gp.get_library_name(organism='Human')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Load data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# load AnnData file\n",
    "data_raw = sc.read_h5ad('adata.h5ad')\n",
    "\n",
    "# remove celltypes that are not present in all subjects\n",
    "data_raw = data_raw[data_raw.obs['celltype']!='End']\n",
    "data_raw = data_raw[data_raw.obs['celltype']!='Per']\n",
    "\n",
    "# display number of cells of each celltype\n",
    "celltypes = data_raw.obs['celltype'].unique().tolist()\n",
    "print(data_raw.obs['celltype'].value_counts())"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Estimate disease pseudo-progression"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# for each celltype, cross-validate to select components and take out-of-sample predictions\n",
    "# look at distribution of these predictions per individual\n",
    "# heatmap of median per celltype and per individual, ordered by PC1\n",
    "subj_col = 'id'\n",
    "ids = data_raw.obs[subj_col].unique().tolist()\n",
    "dx = [data_raw[data_raw.obs[subj_col]==id].obs['diagnosis'].unique()[0] for id in ids]\n",
    "dx = np.array(dx, dtype=int)\n",
    "\n",
    "celltype_preds = {}\n",
    "\n",
    "# whether to perform grid search to select optimal number of PLS components\n",
    "pls_comp_search = False\n",
    "\n",
    "preds = {}\n",
    "for cell in celltypes:\n",
    "    data = data_raw[data_raw.obs['celltype']==cell]\n",
    "    \n",
    "    n_cells_max = 35000\n",
    "    if(data.shape[0]>n_cells_max):\n",
    "        sc.pp.subsample(data, n_obs=n_cells_max, random_state=0)\n",
    "\n",
    "    ### balance classes\n",
    "    data_AD = data[data.obs['diagnosis'] > 0]\n",
    "    data_ctrl = data[data.obs['diagnosis'] <= 0 ]\n",
    "\n",
    "    n_AD = data_AD.shape[0]\n",
    "    n_ctrl = data_ctrl.shape[0]\n",
    "\n",
    "    if(n_AD > n_ctrl):\n",
    "        sc.pp.subsample(data_AD, fraction=n_ctrl/n_AD)\n",
    "    else:\n",
    "        sc.pp.subsample(data_ctrl, fraction=n_AD/n_ctrl)\n",
    "\n",
    "    data = data_AD.concatenate(data_ctrl)\n",
    "    \n",
    "    # remove genes present in less than 0.1% of cells\n",
    "    sc.pp.filter_genes(data, min_cells=int(1+data.shape[0]/1000))\n",
    "    \n",
    "    X = data.X.toarray()\n",
    "    y = data.obs['diagnosis'].values\n",
    "\n",
    "    if(pls_comp_search):\n",
    "        max_comp=8\n",
    "        scPLS = make_pipeline(StandardScaler(), PLSRegression(scale=True))\n",
    "        gridcv = GridSearchCV(estimator=scPLS, param_grid={'plsregression__n_components':[i for i in range(2,max_comp+1)]}, \\\n",
    "            refit=False, scoring='r2', cv=KFold(n_splits=5, shuffle=True, random_state=1), n_jobs=5, return_train_score=False)\n",
    "        _ = gridcv.fit(X, y)\n",
    "        print(f\"{cell}: {gridcv.best_params_['plsregression__n_components']} selected\")\n",
    "        \n",
    "        scPLS_optimal = make_pipeline(StandardScaler(), PLSRegression(n_components=gridcv.best_params_['plsregression__n_components'], scale=False))\n",
    "    else:\n",
    "        # optimal number of components has been previously computed\n",
    "        celltype_n_comps = {'Ex':6, 'Oli':3, 'In':4, 'Ast':2, 'Opc':2, 'Mic':2} # pre-calculated for dataset (clinical diagnosis target)\n",
    "        scPLS_optimal = make_pipeline(StandardScaler(), PLSRegression(n_components=celltype_n_comps[cell], scale=True))\n",
    "    \n",
    "    out = cross_val_predict(scPLS_optimal, X, y, cv=KFold(n_splits=10, shuffle=True, random_state=0), n_jobs=5)\n",
    "\n",
    "    # group predictions by subject id\n",
    "    preds[cell] = []\n",
    "    for id in ids:\n",
    "        id_mask = data.obs[subj_col] == id\n",
    "        if(id_mask.sum()<1):\n",
    "            p = [0]\n",
    "        else:\n",
    "            p = np.squeeze(out[id_mask])\n",
    "\n",
    "        preds[cell].append(p)\n",
    "\n",
    "    celltype_preds[cell] = [np.median(p) for p in preds[cell]]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# create pseudodx\n",
    "pseudodx_df = pd.DataFrame(celltype_preds, index=ids)\n",
    "\n",
    "celltype_dx_pc1 = PCA(n_components=1).fit(pseudodx_df)\n",
    "pseudodx_pc = np.squeeze(celltype_dx_pc1.transform(pseudodx_df))\n",
    "# ensure positive direction corresponds to increasing disease severity\n",
    "pseudodx_pc = pseudodx_pc*np.sign(pseudodx_pc[dx>0].mean())\n",
    "sort_idx = np.argsort(pseudodx_pc, axis=0)\n",
    "ids_sorted = np.array(ids)[sort_idx].tolist()\n",
    "\n",
    "pseudodx_df['dx'] = dx\n",
    "pseudodx_df['dx'] = pseudodx_df['dx'].astype(int)\n",
    "pseudodx_df['sex'] = [float(data_raw[data_raw.obs[subj_col]==id].obs['sex'].unique()[0]=='female') for id in ids]\n",
    "pseudodx_df['pseudodx'] = pseudodx_pc\n",
    "\n",
    "# add ROSMAP metadata from Kellis19 paper\n",
    "pseudodx_df['disease_progression'] = [float(data_raw[data_raw.obs[subj_col]==id].obs['disease_progression'].unique()[0]) for id in ids]\n",
    "pseudodx_df['nft'] = [float(data_raw[data_raw.obs[subj_col]==id].obs['nft'].unique()[0]) for id in ids]\n",
    "pseudodx_df['tangles'] = [float(data_raw[data_raw.obs[subj_col]==id].obs['tangles'].unique()[0]) for id in ids]\n",
    "pseudodx_df['cogn_global_lv'] = [float(data_raw[data_raw.obs[subj_col]==id].obs['cogn_global_lv'].unique()[0]) for id in ids]\n",
    "pseudodx_df['gpath'] = [float(data_raw[data_raw.obs[subj_col]==id].obs['gpath'].unique()[0]) for id in ids]\n",
    "pseudodx_df['amyloid'] = [float(data_raw[data_raw.obs[subj_col]==id].obs['amyloid'].unique()[0]) for id in ids]\n",
    "pseudodx_df['plaq_n'] = [float(data_raw[data_raw.obs[subj_col]==id].obs['plaq_n'].unique()[0]) for id in ids]\n",
    "pseudodx_df['ceradsc'] = [float(data_raw[data_raw.obs[subj_col]==id].obs['ceradsc'].unique()[0]) for id in ids]\n",
    "pseudodx_df['braaksc'] = [float(data_raw[data_raw.obs[subj_col]==id].obs['braaksc'].unique()[0]) for id in ids]\n",
    "\n",
    "pseudodx_df = pseudodx_df.iloc[sort_idx]\n",
    "\n",
    "plt.figure(figsize=(24,8))\n",
    "sns.heatmap(pseudodx_df.loc[:,celltypes[0]:'dx'].T, cmap='coolwarm', vmax=1, vmin=-1, cbar=True)\n",
    "plt.xlabel('Subject ID');\n",
    "\n",
    "save_df = pd.DataFrame(index=pseudodx_df.index, columns=pseudodx_df.columns[:6])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "f, axs = plt.subplots(ncols=len(celltypes), figsize=(20,20), sharey=True)\n",
    "\n",
    "for ax, cell in zip(axs, celltypes):\n",
    "    preds_cell = [preds[cell][i] for i in sort_idx]\n",
    "    ax.vlines(x=0, ymin=0, ymax=len(ids)+1, color='k');\n",
    "    bp = ax.boxplot(preds_cell, vert=False, showfliers=False, patch_artist=True, whis=(10,90))\n",
    "    ax.set_title(f'celltype {cell}')\n",
    "    ax.set_ylim([0,len(ids)+1])\n",
    "    ax.set_yticks(np.arange(1,1+len(ids)))\n",
    "    ax.set_yticklabels(ids)\n",
    "    ax.set_ylabel('subject id')\n",
    "\n",
    "    save_df[cell] = [np.median(preds[cell][i]) for i in sort_idx]\n",
    "\n",
    "    for i, id in enumerate(pseudodx_df.index):\n",
    "        if(int(pseudodx_df['dx'][i])>0):\n",
    "            bp['boxes'][i].set_facecolor('red')\n",
    "        bp['medians'][i].set_color('black')\n",
    "        bp['medians'][i].set_lw(2)\n",
    "\n",
    "id_pseudodx_map = dict(zip(pseudodx_df.index, pseudodx_df['pseudodx']))\n",
    "\n",
    "data_raw.obs['pseudodx'] = data_raw.obs['id'].map(id_pseudodx_map).astype(float)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Verify correlation of disease pseudo-progression with clinical metrics"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "x = ['nft', 'tangles', 'gpath', 'amyloid', 'plaq_n', 'cogn_global_lv', 'braaksc', 'ceradsc']\n",
    "\n",
    "f, axs = plt.subplots(nrows=len(x), figsize=(15,3*len(x)))\n",
    "for ax, metric in zip(axs, x):\n",
    "    dx = pseudodx_df['dx']\n",
    "    ax.plot(pseudodx_df['pseudodx'][dx>0], pseudodx_df[metric][dx>0], 'ro')\n",
    "    ax.plot(pseudodx_df['pseudodx'][dx<=0], pseudodx_df[metric][dx<=0], 'ko')\n",
    "    \n",
    "    percentile = [1,99]\n",
    "    y_smoothed = loess_ci(pseudodx_df['pseudodx'], pseudodx_df[metric], percentile=percentile, frac=2/3)\n",
    "    # plot confidence intervals\n",
    "    for i_ci in range(len(percentile)):\n",
    "        ax.fill_between(pseudodx_df['pseudodx'], y_smoothed[:,1+2*i_ci], y_smoothed[:,2+2*i_ci], alpha=0.5/(i_ci+1), color='g')\n",
    "\n",
    "    ax.set_title(metric)\n",
    "    ax.set_ylabel('principal component 1')\n",
    "pseudodx_df.corr('spearman')['pseudodx']"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# find optimal ordering of subjects given disease severity metrics\n",
    "clinical_metric_pc = PCA(n_components=1).fit_transform(pseudodx_df.loc[:,['nft', 'gpath', 'braaksc', 'ceradsc']])\n",
    "pseudodx_df['clinical_pc'] = -clinical_metric_pc\n",
    "\n",
    "x = ['nft', 'gpath', 'braaksc', 'ceradsc']\n",
    "\n",
    "f, axs = plt.subplots(nrows=len(x), figsize=(15,3*len(x)))\n",
    "for ax, metric in zip(axs, x):\n",
    "    dx = pseudodx_df['dx']\n",
    "    ax.plot(pseudodx_df['clinical_pc'][dx>0], pseudodx_df[metric][dx>0], 'ro')\n",
    "    ax.plot(pseudodx_df['clinical_pc'][dx<=0], pseudodx_df[metric][dx<=0], 'ko')\n",
    "    ax.set_title(metric)\n",
    "pseudodx_df.corr('spearman')['clinical_pc']\n",
    "\n",
    "id_pseudodx_map = dict(zip(pseudodx_df.index, pseudodx_df['clinical_pc']))\n",
    "\n",
    "data_raw.obs['clinical_pc'] = data_raw.obs['id'].map(id_pseudodx_map).astype(float)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Fit PLS-DA model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def filter_shuffle_data(adata_raw, genes, group, subgroup):\n",
    "    data = adata_raw[data_raw.obs[group]==subgroup]\n",
    "\n",
    "    sc.pp.filter_genes(data, min_cells=int(1+data.shape[0]/1000))\n",
    "\n",
    "    genes[subgroup] = np.array(data.var.index.tolist())\n",
    "\n",
    "    # balance classes\n",
    "    binary_target = 'diagnosis'\n",
    "    data_AD = data[data.obs[binary_target] > 0]\n",
    "    data_ctrl = data[data.obs[binary_target] <= 0 ]\n",
    "\n",
    "    n_AD = data_AD.shape[0]\n",
    "    n_ctrl = data_ctrl.shape[0]\n",
    "\n",
    "    if(n_AD > n_ctrl):\n",
    "        sc.pp.subsample(data_AD, fraction=n_ctrl/n_AD, random_state=0)\n",
    "    else:\n",
    "        sc.pp.subsample(data_ctrl, fraction=n_AD/n_ctrl, random_state=0)\n",
    "\n",
    "    data_sub = data_AD.concatenate(data_ctrl)\n",
    "\n",
    "    # shuffle data\n",
    "    data_shuffled = shuffle(data_sub, random_state=0)\n",
    "\n",
    "    return data_shuffled"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# define model parameters\n",
    "group = 'celltype'\n",
    "target = 'diagnosis'\n",
    "scorer = 'roc_auc'\n",
    "celltype_n_comps = {'Mic':2, 'Ast':2, 'Ex':4, 'Oli':3, 'In':3, 'Opc':2} # precomputed\n",
    "\n",
    "loadings_dict = {}\n",
    "gene_symbols = {}\n",
    "scores_df = pd.DataFrame(index=celltypes, columns=[f'r2_fold{i+1}' for i in range(5)]+[f'roc_auc_fold{i+1}' for i in range(5)])\n",
    "\n",
    "for subgroup in ['Mic']:#celltypes:\n",
    "    print(subgroup)\n",
    "    # get data for target celltype\n",
    "    data_shuffled = filter_shuffle_data(data_raw, gene_symbols, group, subgroup)\n",
    "    X = data_shuffled.X.toarray()\n",
    "    y = (data_shuffled.obs[target] > 0).astype(int) # binary\n",
    "\n",
    "    # estimate out-of-sample classification performance\n",
    "    # nested CV\n",
    "    inner_loop = GridSearchCV(PLSRegression(scale=True), param_grid={'n_components':[1,2,3,4,5]})\n",
    "    for metric in ['r2','roc_auc']:\n",
    "        print(f'Metric {metric}')\n",
    "        scores = cross_val_score(inner_loop, X, y, cv=KFold(n_splits=5, shuffle=True), scoring=metric, n_jobs=-1)\n",
    "        print(f'\\t{np.mean(scores):.3f} ({np.std(scores):.3f})')\n",
    "        scores_df.loc[subgroup,f'{metric}_fold1':f'{metric}_fold5'] = scores\n",
    "\n",
    "    # obtain component loadings\n",
    "    optimal_n_comp = celltype_n_comps[subgroup]\n",
    "    scPLS_optimal = PLSRegression(n_components=optimal_n_comp, scale=True)\n",
    "    true_pls = scPLS_optimal.fit(X,y)\n",
    "    true_loadings = true_pls.x_loadings_\n",
    "\n",
    "    loadings_dict[subgroup] = true_loadings\n",
    "\n",
    "    # investigate inhibitory neuron markers\n",
    "    if(subgroup == 'In'):\n",
    "        marker_genes = ['SST','PVALB','KIT','VIP']\n",
    "        x = pd.DataFrame(r, columns=['Module 1', 'Module 2', 'Module 3'], index=marker_genes)\n",
    "        for ic in range(optimal_n_comp):\n",
    "            # generate null for association strength\n",
    "            rs = []\n",
    "            for g in np.random.choice(len(gene_symbols[subgroup]), replace=False, size=1000):\n",
    "                r = spearmanr(true_pls.x_scores_[:,ic], X[:,g])[0]\n",
    "                rs.append(np.abs(r))\n",
    "            print(f'Module {ic+1} null correlation: {np.mean(rs):.2f} ({np.std(rs):.2f} std)')\n",
    "\n",
    "            for g in marker_genes:\n",
    "                r = spearmanr(true_pls.x_scores_[:,ic], X[:,gene_symbols[subgroup]==g])[0]\n",
    "                print(f'\\t{g}: {r:.2f}')\n",
    "\n",
    "                x.loc[g,f'Module {ic+1}'] = r\n",
    "\n",
    "        # create heatmap\n",
    "        plt.figure()\n",
    "        sns.heatmap(x, cmap='bwr', vmin=-0.35, vmax=0.35, annot=True)\n",
    "        plt.show()\n",
    "    \n",
    "    # PHATE visualization\n",
    "    pcs = {'Mic':5, 'Ast':5, 'Opc':5, 'Oli':10, 'In':15, 'Ex':15}\n",
    "    sc.external.tl.phate(data_shuffled, k=15, t='auto', n_pca=pcs[subgroup], a=100, n_jobs=6)\n",
    "\n",
    "    pd.DataFrame(data_shuffled.obsm['X_phate']).to_csv(f'Fig2_{subgroup}_PHATE_data.csv')\n",
    "\n",
    "    for i_comp in range(optimal_n_comp):\n",
    "        data_shuffled.obs['comp_score'] = true_pls.x_scores_[:,i_comp]\n",
    "\n",
    "        lim = max(abs(true_pls.x_scores_[:,i_comp]))\n",
    "        dot_size = 20\n",
    "\n",
    "        # PHATE visualization, colored by module score\n",
    "        with plt.rc_context({'figure.figsize':[6,6]}):\n",
    "            plt.figure()\n",
    "            sc.external.pl.phate(data_shuffled, color='comp_score', cmap='coolwarm', norm=None, \\\n",
    "                                size=dot_size, title=f'Cell scores component {i_comp+1}', show=False)\n",
    "        \n",
    "        # PHATE visualization, colored by diagnosis\n",
    "        with plt.rc_context({'figure.figsize':[6,6]}):\n",
    "            dot_size = 30\n",
    "            plt.figure()\n",
    "            sc.external.pl.phate(data_shuffled, color='diagnosis', cmap='coolwarm', size=dot_size, alpha=0.5, show=False)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "f, ax = plt.subplots()\n",
    "for i1,c in enumerate(celltypes):\n",
    "    for i2,metrics in enumerate(['r2','roc_auc']):\n",
    "        ax.plot(i2+2*np.array(5*[i1]),scores_df.loc[c,f'{metric}_fold1':f'{metric}_fold5'], '.')\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# plot of out-of-sample classification performance (Figure 1b)\n",
    "f, ax = plt.subplots()\n",
    "# ax.boxplot(scores_df.loc[:,'r2_fold1':'r2_fold5'].T, positions=np.arange(0,15,2.5), showfliers=False, showbox=False, widths=0.8, whis=(1,99));\n",
    "for i1,c in enumerate(celltypes):\n",
    "    ax.plot(2.5*np.array(5*[i1]),scores_df.loc[c,'r2_fold1':'r2_fold5'], 'ok', ms=10)\n",
    "ax2 = ax.twinx()\n",
    "# ax2.boxplot(scores_df.loc[:,'roc_auc_fold1':'roc_auc_fold5'].T, positions=1+np.arange(0,15,2.5), showfliers=False, showbox=False, widths=0.8, whis=(1,99));\n",
    "for i1,c in enumerate(celltypes):\n",
    "    ax2.plot(1+2.5*np.array(5*[i1]),scores_df.loc[c,'roc_auc_fold1':'roc_auc_fold5'], 'ok', ms=10)\n",
    "\n",
    "ax.set_xticks(0.5+np.arange(0,15,2.5), labels=celltypes);\n",
    "plt.show()"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Create heatmap of top modules genes"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "group = 'celltype'\n",
    "target = 'diagnosis'\n",
    "scorer = 'roc_auc'\n",
    "\n",
    "celltype_n_comps = {'Mic':2, 'Ast':2, 'Ex':4, 'Oli':3, 'In':3, 'Opc':2}\n",
    "module_correlations = {k:[] for k in celltype_n_comps.keys()}\n",
    "gene_symbols = {}\n",
    "loadings = {}\n",
    "\n",
    "for subgroup in celltype_n_comps.keys():\n",
    "    data_shuffled = filter_shuffle_data(data_raw, gene_symbols, group, subgroup)\n",
    "    X = data_shuffled.X.toarray()\n",
    "    y = (data_shuffled.obs[target] > 0).astype(int) # binary\n",
    "\n",
    "    # fit model\n",
    "    optimal_n_comp = celltype_n_comps[subgroup]\n",
    "    scPLS_optimal = PLSRegression(n_components=optimal_n_comp, scale=True)\n",
    "    true_pls = scPLS_optimal.fit(X,y)\n",
    "    true_scores = true_pls.x_scores_\n",
    "\n",
    "    # compute correlations\n",
    "    for i_comp in range(optimal_n_comp):\n",
    "        module_correlations[subgroup].append(np.corrcoef(true_scores[:,i_comp], y)[0,1])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# heatmap of correlations between scores and diagnosis\n",
    "fill = np.empty((6,4))\n",
    "fill[:] = np.nan\n",
    "corrs = pd.DataFrame(fill, index=module_correlations.keys(), columns=[f'Module {i+1}' for i in range(4)])\n",
    "for subgroup in module_correlations.keys():\n",
    "    for i_comp in range(len(module_correlations[subgroup])):\n",
    "        corrs.loc[subgroup, f'Module {i_comp+1}'] = np.abs(module_correlations[subgroup][i_comp])\n",
    "\n",
    "plt.figure()\n",
    "sns.heatmap(corrs, cmap='vlag', annot=True, vmin=-0.70, vmax=0.70)\n",
    "plt.xticks(rotation=45, ha='right');\n",
    "plt.yticks(rotation=0);\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# heatmap of top genes\n",
    "top_n = 3\n",
    "modules = [f'{k}{i+1}' for (k,v) in celltype_n_comps.items() for i in range(v)]\n",
    "n_modules = len(modules)\n",
    "\n",
    "# generate index\n",
    "index_tmp = []\n",
    "for c in gene_symbols.keys():\n",
    "    for m in range(celltype_n_comps[c]):\n",
    "        index_tmp.extend(np.array(gene_symbols[c])[np.argsort(loadings[c][:,m])[-top_n:]][::-1])\n",
    "\n",
    "top_gene_loadings = pd.DataFrame(np.zeros((len(index_tmp), n_modules)), columns=modules, index=index_tmp)\n",
    "\n",
    "# fill in loadings\n",
    "for c in gene_symbols.keys(): # celltypes\n",
    "    for m in range(celltype_n_comps[c]): # modules\n",
    "        for g in top_gene_loadings.index:\n",
    "            if g in gene_symbols[c]:\n",
    "                top_gene_loadings.loc[g,f'{c}{m+1}'] = loadings[c][gene_symbols[c]==g,m]\n",
    "\n",
    "top_gene_loadings.drop_duplicates(inplace=True)\n",
    "\n",
    "plt.figure(figsize=(10,top_gene_loadings.shape[0]*0.2))\n",
    "sns.heatmap(top_gene_loadings, cmap='vlag', vmin=-0.07, vmax=0.07)\n",
    "plt.show()"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Run bootstrap analysis (per celltype)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "subgroup = 'Mic'\n",
    "true_loadings = loadings_dict[subgroup]\n",
    "\n",
    "### perturbation (get bootstrap distribution of loadings)\n",
    "data_shuffled = filter_shuffle_data(data_raw, gene_symbols, group, subgroup)\n",
    "\n",
    "n_bootstrap = 1000\n",
    "def parallel_bootstrap(j, data_bs):\n",
    "    if((100*j/n_bootstrap)%10==0):\n",
    "        print(f'{100*j/n_bootstrap:.0f}% complete')\n",
    "    rng = default_rng(j)\n",
    "\n",
    "    # select bootstrap data\n",
    "    idx_bootstrap = []\n",
    "    for i in data_bs.obs[target].unique():\n",
    "        idx_class = np.where(data_bs.obs[target] == i)[0]\n",
    "        n_class = idx_class.shape[0]\n",
    "        idx_class_bootstrap = rng.integers(0, n_class, n_class)\n",
    "        idx_bootstrap.extend(idx_class[idx_class_bootstrap])\n",
    "    \n",
    "    data_bootstrap = data_bs[idx_bootstrap, :]\n",
    "    data_bootstrap.obs_names_make_unique()\n",
    "    \n",
    "    X_bootstrap = data_bootstrap.X.toarray()\n",
    "    y_bootstrap = (data_bootstrap.obs[target] > 0).astype(int)\n",
    "\n",
    "    # calculate model loadings of null model\n",
    "    bootstrap_loadings = scPLS_optimal.fit(X_bootstrap, y_bootstrap).x_loadings_\n",
    "\n",
    "    return bootstrap_loadings\n",
    "\n",
    "bootstrap_loadings = Parallel(n_jobs=4)(delayed(parallel_bootstrap)(j, data_shuffled) for j in range(n_bootstrap))\n",
    "bootstrap_loadings = np.array(bootstrap_loadings)\n",
    "\n",
    "# bootstrap_loadings shape: [n_bootstrap, n_genes, n_components]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# some bootstrap loadings are mirrored, which causes problems with distribution\n",
    "# for each bootstrap loading, determine whether flipping it results in a lower distance between it and the true loading\n",
    "mirror = True\n",
    "if(mirror):\n",
    "    mirror_mask = np.mean(np.abs(bootstrap_loadings - true_loadings), axis=1) > np.mean(np.abs(-bootstrap_loadings - true_loadings), axis=1)\n",
    "    #mirror_mask = correlation\n",
    "    mirror_factor = np.where(mirror_mask, -1, 1)\n",
    "\n",
    "    mirror_factor = np.expand_dims(mirror_factor, axis=1)\n",
    "\n",
    "    bootstrap_loadings = bootstrap_loadings * mirror_factor\n",
    "\n",
    "# calculate mean of bootstrap distribution for each feature\n",
    "bootstrap_mean = np.median(bootstrap_loadings, axis=0)\n",
    "\n",
    "# create copies to modify\n",
    "bootstrap_mean_zeroed = bootstrap_mean.copy()\n",
    "bootstrap_loadings_zeroed = bootstrap_loadings.copy()\n",
    "\n",
    "# zero distributions that significantly cross zero\n",
    "zero_threshold = 5 # fraction of bootstrap loadings that need to cross zero to zero out feature\n",
    "limits = np.percentile(bootstrap_loadings_zeroed, q=[zero_threshold,100-zero_threshold], axis=0)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "gene_symbols"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# boolean mask for features where the sign of one of the limits is of opposite sign from the median\n",
    "# True indicates feature should be dropped\n",
    "zero_mask = np.logical_or(np.logical_xor(limits[0,:,:]>0, bootstrap_mean>0), np.logical_xor(limits[1,:,:]>0, bootstrap_mean>0))\n",
    "zero_mask = np.logical_or(zero_mask, (bootstrap_loadings_zeroed==0).sum(axis=0)>0)\n",
    "\n",
    "bootstrap_mean_zeroed[zero_mask] = 0\n",
    "bootstrap_loadings_zeroed[:,zero_mask] = 0\n",
    "# bootstrap_mean = np.where(zero_mask, 0, bootstrap_mean)\n",
    "\n",
    "for i in range(zero_mask.shape[1]):\n",
    "    print(f\"component {i}: {zero_mask[:,i].sum()} features zeroed\")\n",
    "\n",
    "# save component loadings for later use\n",
    "df = pd.DataFrame(bootstrap_mean_zeroed, index=gene_symbols[subgroup])\n",
    "df.to_csv(f'data/{subgroup}_significant_loadings.csv')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Compare PLS modules to DEGs"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# correlated PLS module gene weights with DEG results\n",
    "degs = {}\n",
    "loadings = {}\n",
    "\n",
    "fill = np.empty((4,6))\n",
    "fill[:] = np.nan\n",
    "heatmap = pd.DataFrame(fill, index=[f'Module {i}' for i in range(1,5)], columns=celltypes)\n",
    "for c in ['Ex','In','Oli','Mic','Ast','Opc']:\n",
    "    degs[c] = pd.read_excel('data/41586_2019_1195_MOESM4_ESM.xlsx', sheet_name=c, index_col=0)\n",
    "    loadings[c] = pd.read_csv(f'data/{c}_significant_loadings.csv', index_col=0)\n",
    "\n",
    "    # get genes present in both sets\n",
    "    shared_genes = list(set(degs[c].index).intersection(set(loadings[c].index)))\n",
    "\n",
    "    deg_feature = 'IndModel.FC' # which metric to use for comparison\n",
    "    for comp in range(loadings[c].shape[1]):\n",
    "        comp_loading = loadings[c].loc[shared_genes].iloc[:,comp].values # sorted by shared genes\n",
    "        deg_metric = degs[c].loc[shared_genes, deg_feature].values\n",
    "        deg_metric[np.isnan(deg_metric)] = 0\n",
    "        r,p = pearsonr(comp_loading, deg_metric)\n",
    "        \n",
    "        heatmap.loc[f'Module {comp+1}', c] = r\n",
    "\n",
    "plt.figure(constrained_layout=True)\n",
    "sns.heatmap(heatmap.T, cmap='vlag', annot=True, vmin=-1, vmax=1)\n",
    "plt.xticks(rotation=45, ha='right');\n",
    "plt.yticks(rotation=0);\n",
    "plt.ylabel('Celltype')\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# look at top DEGs by p-value and FC\n",
    "for c in ['Ex','In','Oli','Mic','Ast','Opc']:\n",
    "    DEG = pd.read_excel('data/41586_2019_1195_MOESM4_ESM.xlsx', sheet_name=c, index_col=0)\n",
    "    PLS = pd.read_csv(f'data/{c}_significant_loadings.csv', index_col=0)\n",
    "\n",
    "    deg_feature = 'IndModel.adj.pvals' #'IndModel.FC' # which metric to use for comparison\n",
    "    for comp in range(PLS.shape[1]):\n",
    "        pls_genes_sorted = PLS.abs().sort_values(by=f'{comp}', ascending=False).index.tolist()[:10]\n",
    "        deg_genes_sorted = DEG.abs().sort_values(by=deg_feature, ascending=True).index.tolist()[:10]\n",
    "        \n",
    "        gene_overlap = list(set(pls_genes_sorted).intersection(set(deg_genes_sorted)))\n",
    "        print(f'{c} module {comp+1}: {gene_overlap}')"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Plot top genes per module"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# plot significant features\n",
    "f, axs = plt.subplots(nrows=optimal_n_comp, figsize=(25,optimal_n_comp*15))\n",
    "if(optimal_n_comp==1):\n",
    "    axs = [axs]\n",
    "\n",
    "N_TOP = 15\n",
    "\n",
    "# plot the significant genes for each component\n",
    "for comp, ax in enumerate(axs):\n",
    "    sort_idx = np.argsort(np.abs(bootstrap_mean_zeroed[:,comp]))[::-1]\n",
    "\n",
    "    # plot top genes\n",
    "    NN = 0\n",
    "    ax.violinplot(bootstrap_loadings_zeroed[:,:,comp][:,sort_idx][:,NN:(NN+N_TOP)], showextrema=False, widths=0.8, showmedians=True, points=200);\n",
    "    ax.set_xticks(np.arange(1,N_TOP+1))\n",
    "    ax.set_xticklabels(labels=gene_symbols[subgroup][sort_idx][NN:(NN+N_TOP)], rotation=65, fontsize=30)\n",
    "    ax.set_xlim([0,N_TOP+1])\n",
    "plt.show()"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Obtain GSEA enriched pathways per module"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "gene_set_list = ['GO_Biological_Process_2021', 'WikiPathway_2021_Human', 'Panther_2016']\n",
    "gsea_terms_tmp = []\n",
    "# perform GSEA for each component in the model\n",
    "for comp in range(optimal_n_comp):\n",
    "    gsea_terms_tmp_comp = pd.DataFrame()\n",
    "    # perform GSEA using multiple gene set databases\n",
    "    for gene_set in gene_set_list:\n",
    "        ranked_genes = pd.DataFrame(data={'genes':gene_symbols[subgroup][~zero_mask[:,comp]].tolist(), 'PLS_weights':bootstrap_mean[:,comp][~zero_mask[:,comp]]})\n",
    "        try:\n",
    "            pre_res = gp.prerank(rnk=ranked_genes,\n",
    "                            gene_sets=gene_set,\n",
    "                            processes=4,\n",
    "                            min_size=5,\n",
    "                            no_plot=True,\n",
    "                            permutation_num=1000,\n",
    "                            outdir=None,\n",
    "                            seed=1)\n",
    "        except Exception:\n",
    "            print(f'GSEA Error: No enriched gene modules found in component {comp}')\n",
    "        \n",
    "        df = pre_res.res2d\n",
    "        df['component'] = comp\n",
    "        df['gene_set_source'] = gene_set\n",
    "        df = df[df['fdr']<0.05]\n",
    "        df = df[~df['pval'].isna()]\n",
    "        df = df.sort_values('nes', ascending=False)\n",
    "        display(df.head(10))\n",
    "        gsea_terms_tmp_comp = pd.concat((gsea_terms_tmp_comp, df), axis=0)\n",
    "    gsea_terms_tmp.append(gsea_terms_tmp_comp)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# process terms for each component, then combine terms from all components into one dataframe\n",
    "gsea_terms = pd.DataFrame()\n",
    "for t in gsea_terms_tmp:\n",
    "    # remove duplicated terms\n",
    "    t = t[~t.index.duplicated()]\n",
    "    t.index = t.index.values + str(t['component'][0])\n",
    "\n",
    "    # sort by absolute value of gsea score\n",
    "    t = t.reindex(t['nes'].abs().sort_values(ascending=False).index)\n",
    "\n",
    "    gsea_term_sets = [set(x.split(';')) for x in t['genes'].values]\n",
    "\n",
    "    # list to track number of times a term is a superset of another\n",
    "    superset_count_list = [0]*len(gsea_term_sets)\n",
    "\n",
    "    # drop terms that are superset of others\n",
    "    set_overlap = []\n",
    "    drop_idx = []\n",
    "    drop_bool = False\n",
    "    overlap_threshold = 0.9\n",
    "\n",
    "    # select a gene set\n",
    "    for i in range(len(gsea_term_sets)):\n",
    "        # skip terms that will be dropped\n",
    "        if((i in drop_idx) and (drop_bool)):\n",
    "            continue\n",
    "        \n",
    "        # check if selected gene set overlaps with any other remaining gene set\n",
    "        # drop less enriched (lower enrichment score) gene set if overlap is above threshold\n",
    "        tmp_overlap = []\n",
    "        for j in range(i+1, len(gsea_term_sets)):\n",
    "            # size of smaller set\n",
    "            min_size = min(len(gsea_term_sets[i]), len(gsea_term_sets[j]))\n",
    "            # number of genes common between sets\n",
    "            overlap = len(gsea_term_sets[i].intersection(gsea_term_sets[j]))\n",
    "            # save index if overlap is too high\n",
    "            \n",
    "            if(overlap/min_size >= overlap_threshold):\n",
    "                drop_idx.append(j)\n",
    "                superset_count_list[j] += 1\n",
    "\n",
    "    # drop terms with high overlap\n",
    "    if(drop_bool):\n",
    "        keep_idx = [i for i in range(len(gsea_term_sets)) if i not in drop_idx]\n",
    "        t = t.iloc[keep_idx]\n",
    "\n",
    "        print(f'dropping {len(drop_idx)} terms, {len(keep_idx)} remain')\n",
    "\n",
    "    gsea_terms = pd.concat((gsea_terms, pd.concat((t['nes'], t['fdr'], t['component'], t['gene_set_source'], t['genes'].str.split(';'), t['ledge_genes'].str.split(';'), pd.Series(data=superset_count_list, index=t.index)), axis=1)), axis=0)\n",
    "\n",
    "gsea_terms.rename(columns={0:'superset_count'}, inplace=True)\n",
    "\n",
    "print(gsea_terms.shape)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# identify AD GWAS genes in leading edge genes\n",
    "# gene list from: Wightman, D. P. et al. A genome-wide association study \n",
    "# with 1,126,563 individuals identifies new risk loci for Alzheimer’s disease. Nat. Genet. 53, 1276–1282 (2021).\n",
    "goi = {\n",
    "    'AD GWAS genes':set(['AGRN','CR1','NCK2','BIN1','INPPD5','CLNK','TNIP1','HAVCR2','HLA-DRB1','TREM2','CD2AP',\n",
    "    'TMEM106B','ZCWPW1', 'NYAP1','EPHA1-AS1','CLU','SHARPIN','USP6NL', 'ECHDC3','CCDC6','MADD', 'SPI1','MS4A4A','PICALM',\n",
    "    'SORL1','FERMT2','RIN3','ADAM10','APH1B','SCIMP', 'RABEP1','GRN','ABI3','TSPOAP1-AS1','ACE','ABCA7','APOE','NTN5',\n",
    "    'CD33','LILRB2','CASS4','APP', 'IFNB1']),\n",
    "    }\n",
    "\n",
    "def identify_genes(row, gene_list):\n",
    "    hits = []\n",
    "    for g in gene_list:\n",
    "        if(g in row['genes']):\n",
    "            hits.append(g)\n",
    "    \n",
    "    return hits\n",
    "\n",
    "for k,v in goi.items():\n",
    "    gsea_terms[k] = gsea_terms.apply(identify_genes, args=(v,), axis=1)"
   ]
  }
 ],
 "metadata": {
  "interpreter": {
   "hash": "c5169af85f5023cc4c2fea288b97483cc61fa2093ddd8ea31e42060c7c504a03"
  },
  "kernelspec": {
   "display_name": "Python 3.9.7 ('base')",
   "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.9.16"
  },
  "orig_nbformat": 4
 },
 "nbformat": 4,
 "nbformat_minor": 2
}