{
"cells": [
{
"cell_type": "markdown",
"metadata": {},
"source": [
"# Batch Correction Tutorial\n",
"\n",
"Authors: [Hui Ma](https://github.com/huimalinda), [Yiming Yang](https://github.com/yihming), [Rimte Rocher](https://github.com/rocherr)
\n",
"Date: 2022-04-10
\n",
"Notebook Source: [batch_correction.ipynb](https://raw.githubusercontent.com/lilab-bcb/pegasus-tutorials/main/notebooks/batch_correction.ipynb)"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"import pegasus as pg"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"## Test Environment\n",
"\n",
"Some batch correction methods allow use GPU for acceleration, which we will enable below. Below is a summary on the test environment when compiling this notebook:\n",
"\n",
"|CPU Model|CPU(s)|GPU Model|GPU Memory|GPU Driver version|CUDA version|Operating System|\n",
"|:---:|:---:|:---:|:---:|:---:|:---:|:---:|\n",
"|AMD Ryzen 7 2700X|16|NVIDIA GeForce 1050 Ti|4GB|470.103.01|11.4|Ubuntu Linux 21.10|"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"## Dataset"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"In this tutorial, we'll use a gene-count matrix dataset on human bone marrow from 8 donors, and show how to use the batch correction methods in Pegasus to tackle the batch effects in the data.\n",
"\n",
"The dataset is stored at https://storage.googleapis.com/terra-featured-workspaces/Cumulus/MantonBM_nonmix_subset.zarr.zip. You can also use [gsutil](https://cloud.google.com/storage/docs/gsutil) to download it via its Google bucket URL (`gs://terra-featured-workspaces/Cumulus/MantonBM_nonmix_subset.zarr.zip`)."
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"Now load the count matrix:"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"data = pg.read_input(\"MantonBM_nonmix_subset.zarr.zip\")\n",
"data"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"`'Channel'` is the batch key. Each batch is associated with one donor, so there are 8 batches in total."
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"## Sections\n",
"- [Preprocessing](#pre)\n",
"- [Clustering without Correcting Batch Effects](#nobatch)\n",
"- [Batch Correction Methods](#batch)\n",
" - [Harmony](#harmony)\n",
" - [Integrative-NMF](#inmf)\n",
" - [scVI](#scvi)\n",
" - [Scanorama](#scanorama)\n",
"- [Comparing Batch CorrectionMethods](#comp)"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"## Preprocessing\n",
"\n",
"First, preprocess the data. This includes filtration, selecting robust genes, and log-normalization:"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"pg.qc_metrics(data, min_genes=500, max_genes=6000, mito_prefix='MT-', percent_mito=10)\n",
"pg.filter_data(data)\n",
"pg.identify_robust_genes(data)\n",
"pg.log_norm(data)"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"After quality-control, distribution of cells in all the 8 batches is:"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"data.obs['Channel'].value_counts()"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"## Clustering without Correcting Batch Effects\n",
"\n",
"We first perform downstream steps without considering batch effects. In this way, you can see where the batch effects exist, and moreover, we'll use this result as the baseline when comparing different batch correction methods."
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"data_baseline = data.copy()\n",
"pg.highly_variable_features(data_baseline)"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"In this tutorial, the downstream steps consists of:\n",
"* PCA calculation: by default calculate 50 PCs;\n",
"* kNN graph construction, and Louvain clustering based on it;\n",
"* Calculate 2-dimensinoal UMAP embedding, and show UMAP plot."
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"pg.pca(data_baseline)\n",
"pg.neighbors(data_baseline)\n",
"pg.louvain(data_baseline)\n",
"pg.umap(data_baseline)\n",
"pg.scatter(data_baseline, attrs=['louvain_labels', 'Channel'], basis='umap')"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"Let's have a quick look at the UMAP plot above. If checking the cells in Louvain cluster 11 and 14 in the right-hand-side plot, you can see that most of them are from sample `MantonBM3_HiSeq_1`. This indicates strong batch effects."
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"## Batch Correction Methods\n",
"\n",
"Batch effect occurs when data samples are generated in different conditions, such as date, weather, lab setting, equipment, etc. Unless informed that all the samples were generated under the similar condition, people may suspect presumable batch effects if they see a visualization graph with samples kind-of isolated from each other.\n",
"\n",
"In this tutorial, you'll see how to apply the batch correction methods in Pegasus to this dataset. \n",
"\n",
"As a common step ahead, we need to re-select HVGs considering batch effects:"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"pg.highly_variable_features(data, batch='Channel')"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"### Harmony Algorithm\n",
"\n",
"[Harmony](https://www.nature.com/articles/s41592-019-0619-0) is a widely-used method for data integration. Pegasus uses [harmony-pytorch](https://github.com/lilab-bcb/harmony-pytorch) package to perform Harmony batch correction.\n",
"\n",
"Harmony works on PCA matrix. So we need to first calculate the original PCA matrix:"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"data_harmony = data.copy()\n",
"pg.pca(data_harmony)"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"Now we are ready to run Harmony integration:"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"harmony_key = pg.run_harmony(data_harmony, use_gpu=True)"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"When finished, the corrected PCA matrix is stored in `data_harmony.obsm['X_pca_harmony']`, and `run_harmony` returns the representation key `'pca_harmony'` as variable `harmony_key`. In the downstream steps, you can set `rep` parameter to either `harmony_key` or `'pca_harmony'` in Pegasus functions whenever applicable.\n",
"\n",
"Notice that by default `run_harmony` has `use_gpu=False` which uses only CPUs. In this notebook, as above, we can set this parameter to `True` to enable GPU computation to speed up.\n",
"\n",
"For details on parameters of `run_harmony` other than the default setting, please [see here](https://pegasus.readthedocs.io/en/stable/api/pegasus.run_harmony.html#pegasus.run_harmony).\n",
"\n",
"With the new corrected PCA matrix, we can perform kNN-graph-based clustering and calculate UMAP embeddings as follows:"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"pg.neighbors(data_harmony, rep=harmony_key)\n",
"pg.louvain(data_harmony, rep=harmony_key)\n",
"pg.umap(data_harmony, rep=harmony_key)"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"Then show UMAP plot:"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"pg.scatter(data_harmony, attrs=['louvain_labels', 'Channel'], basis='umap')"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"### Integrative-NMF\n",
"\n",
"Another popular data integration method is [integrative Non-negative Matrix Factorization](https://academic.oup.com/bioinformatics/article/32/1/1/1743821) (iNMF). Pegasus includes this method in [nmf-torch](https://github.com/lilab-bcb/nmf-torch) package."
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"data_inmf = data.copy()\n",
"inmf_key = pg.integrative_nmf(data_inmf, batch='Channel', use_gpu=True)"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"Similarly as `run_harmony`, when finished, the calculated embedding is stored in `data_inmf.obsm['X_inmf']`, and `integrative_nmf` returns the representation key `'inmf'` as variable `inmf_key`. In the downstream steps, you can set `rep` parameter to either `inmf_key` or `'inmf'` in Pegasus functions whenever applicable.\n",
"\n",
"Notice that by default `integrative_nmf` has `use_gpu=False` which uses only CPUs. In this notebook, as above, we can set this parameter to `True` to enable GPU computation to speed up.\n",
"\n",
"For details on parameters of `integrative_nmf` other than the default setting, please [see here](https://pegasus.readthedocs.io/en/stable/api/pegasus.integrative_nmf.html#pegasus.integrative_nmf).\n",
"\n",
"Now we can perform kNN-graph-based clustering and calculate UMAP embeddings as follows:"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"pg.neighbors(data_inmf, rep=inmf_key)\n",
"pg.louvain(data_inmf, rep=inmf_key)\n",
"pg.umap(data_inmf, rep=inmf_key)"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"Then show UMAP plot:"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"pg.scatter(data_inmf, attrs=['louvain_labels', 'Channel'], basis='umap')"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"### scVI\n",
"\n",
"[scVI](https://www.nature.com/articles/s41587-021-01206-w) is a popular algorithm for data integration. Pegasus uses [scVI-tools](https://github.com/scverse/scvi-tools) package.\n",
"\n",
"scVI uses a Deep Learning model to learn embeddings directly from the **raw counts** after HVG selection:"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"data_scvi = data.copy()\n",
"scvi_key = pg.run_scvi(data_scvi, batch='Channel', use_gpu=True)"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"When finished, the calculated embedding is stored in `data_scvi.obsm['X_scVI']`, and `run_scvi` returns the representation key `'scVI'` as variable `scvi_key`. In the downstream steps, you can set `rep` parameter to either `scvi_key` or `'scVI'` in Pegasus functions whenever applicable.\n",
"\n",
"Notice that by default `run_scvi` has `use_gpu=False` which uses only CPUs. In this notebook, as above, we can set this parameter to `True` to enable GPU computation to speed up.\n",
"\n",
"For details on parameters of `run_scvi` other than the default setting, please [see here](https://pegasus.readthedocs.io/en/stable/api/pegasus.run_scvi.html#pegasus.run_scvi).\n",
"\n",
"Now we can perform kNN-graph-based clustering and calculate UMAP embeddings as follows:"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"pg.neighbors(data_scvi, rep=scvi_key)\n",
"pg.louvain(data_scvi, rep=scvi_key)\n",
"pg.umap(data_scvi, rep=scvi_key)"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"Then show UMAP plot:"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"pg.scatter(data_scvi, attrs=['louvain_labels', 'Channel'], basis='umap')"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"### Scanorama Algorithm\n",
"\n",
"[Scanorama](https://www.nature.com/articles/s41587-019-0113-3) is another widely-used batch correction algorithm. Pegasus uses [scanorama](https://github.com/brianhie/scanorama) package.\n",
"\n",
"Similarly as Harmony, Scanorama corrects the original PCA matrix of the dataset. But as PCA step is already integrated in `run_scanorama` function, we don't need to run `pca` seperately:"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"data_scan = data.copy()\n",
"scan_key = pg.run_scanorama(data_scan)"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"You can check details on `run_scanorama` parameters [here](https://pegasus.readthedocs.io/en/stable/api/pegasus.run_scanorama.html).\n",
"\n",
"By default, it considers count matrix only regarding the selected HVGs, and calculates the corrected PCA matrix of $50$ PCs. When finished, this new PCA matrix is stored in `data_scan.obsm['X_scanorama']`, and returns its representation key `'scanorama'` as variable `scan_key`. In the downstream steps, you can set `rep` parameter to either `scan_key` or `'scanorama'` in Pegasus functions whenever applicable:"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"pg.neighbors(data_scan, rep=scan_key)\n",
"pg.louvain(data_scan, rep=scan_key)\n",
"pg.umap(data_scan, rep=scan_key)"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"Now check its UMAP plot:"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"pg.scatter(data_scan, attrs=['louvain_labels', 'Channel'], basis='umap')"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"## Comparing Batch Correction Methods\n",
"\n",
"\n",
"To compare the performance on the three methods, one metric is **runtime**, which you can see from the logs in sections above: integrativ nmf method is the fastest, then Harmony, and Scanorama is the slowest.\n",
"\n",
"In this section, we'll use 2 other metrics for comparison:\n",
"* **kBET acceptance rate:** kBET measures whether batches are well-mixed in the local neighborhood of each cell. Then kBET acceptance rate is the percent of cells with kBET p-values $\\geq 0.05$. The higher, the better.\n",
"* **kSIM acceptance rate:** kSIM measures whether cells of the same pre-annotated cell type are still close to each other in the local neighborhoods after batch correction. Then kSIM acceptance rate is the percent of cells with kSIM p-values $\\geq 0.05$. We use this metric to reflect whether known biological relationships are preserved after correction. The higher, the better.\n",
"\n",
"We have 4 results: No batch correction (Baseline), Harmony, iNMF, and Scanorama. For each result, kBET and kSIM acceptance rates are calculated on its 2D UMAP coordinates for comparison, which is consistent with [Cumulus paper](https://rdcu.be/b5R5B).\n",
"\n",
"Details on these 2 metrics can also be found in Cumulus paper."
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"### kBET Acceptance Rate\n",
"\n",
"We can use `calc_kBET` function to calculate on kBET acceptance rates. Besides,\n",
"* For `attr` parameter, use the batch key, which is `'Channel'` in this tutorial.\n",
"* For `rep` parameter, set to the corresponding UMAP coordinates;\n",
"* For returned values, we only care about kBET acceptance rates, so just ignore the first two returned values."
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"_, _, kBET_baseline = pg.calc_kBET(data_baseline, attr='Channel', rep='umap')\n",
"_, _, kBET_harmony = pg.calc_kBET(data_harmony, attr='Channel', rep='umap')\n",
"_, _, kBET_inmf = pg.calc_kBET(data_inmf, attr='Channel', rep='umap')\n",
"_, _, kBET_scvi = pg.calc_kBET(data_scvi, attr='Channel', rep='umap')\n",
"_, _, kBET_scan = pg.calc_kBET(data_scan, attr='Channel', rep='umap')"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"### kSIM Acceptance Rate\n",
"\n",
"We need pre-annotated cell type information as ground truth to calculate kSIM acceptance rate. This is achieved by:\n",
"* Starting with Louvain clusters in the baseline case;\n",
"* Merging two CD14+ Monocyte clusters, as one is donor-specific;\n",
"* For the only cluster with no strong cell type evidence, classifying its cells using a Gradient Boosting model trained from its neighbor T-cell clusters.\n",
"\n",
"This ground truth is stored at https://storage.googleapis.com/terra-featured-workspaces/Cumulus/cell_types.csv, or you can get it using `gsutil` from its Google bucket URL (gs://terra-featured-workspaces/Cumulus/cell_types.csv):"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"Now load this file, and attach its `'cell_types'` column to the 4 resulting count matrices above:"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"import pandas as pd\n",
"import numpy as np\n",
"\n",
"df_celltypes = pd.read_csv(\"cell_types.csv\", index_col='barcodekey')\n",
"\n",
"assert np.sum(df_celltypes.index!=data_baseline.obs_names) == 0\n",
"data_baseline.obs['cell_types'] = df_celltypes['cell_types']\n",
"\n",
"assert np.sum(df_celltypes.index!=data_harmony.obs_names) == 0\n",
"data_harmony.obs['cell_types'] = df_celltypes['cell_types']\n",
"\n",
"assert np.sum(df_celltypes.index!=data_inmf.obs_names) == 0\n",
"data_inmf.obs['cell_types'] = df_celltypes['cell_types']\n",
"\n",
"assert np.sum(df_celltypes.index!=data_scvi.obs_names) == 0\n",
"data_scvi.obs['cell_types'] = df_celltypes['cell_types']\n",
"\n",
"assert np.sum(df_celltypes.index!=data_scan.obs_names) == 0\n",
"data_scan.obs['cell_types'] = df_celltypes['cell_types']"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"We can then use `calc_kSIM` function to calculate kSIM acceptance rates. Besides,\n",
"* For `attr` parameter, use the ground truth key `'cell_types'`;\n",
"* For `rep` parameter, similarly as in kBET section, set to the corresponding UMAP coordinates;\n",
"* For returned values, we only care about kSIM acceptance rates, so just ignore the first returned values."
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"_, kSIM_baseline = pg.calc_kSIM(data_baseline, attr='cell_types', rep='umap')\n",
"_, kSIM_harmony = pg.calc_kSIM(data_harmony, attr='cell_types', rep='umap')\n",
"_, kSIM_inmf = pg.calc_kSIM(data_inmf, attr='cell_types', rep='umap')\n",
"_, kSIM_scvi = pg.calc_kSIM(data_scvi, attr='cell_types', rep='umap')\n",
"_, kSIM_scan = pg.calc_kSIM(data_scan, attr='cell_types', rep='umap')"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"### Performance Plot\n",
"\n",
"Now draw a scatterplot regarding these two metrics on the results:"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"import seaborn as sns\n",
"import matplotlib.pyplot as plt\n",
"\n",
"df_plot = pd.DataFrame({'method': ['Baseline', 'Harmony', 'iNMF', 'scVI', 'Scanorama'],\n",
" 'kBET': [kBET_baseline, kBET_harmony, kBET_inmf, kBET_scvi, kBET_scan],\n",
" 'kSIM': [kSIM_baseline, kSIM_harmony, kSIM_inmf, kSIM_scvi, kSIM_scan]})\n",
"\n",
"plt.figure(dpi=100)\n",
"ax = sns.scatterplot(x = 'kSIM', y = 'kBET', hue = 'method', data = df_plot, legend = False)\n",
"for line in range(0, df_plot.shape[0]):\n",
" x_pos = df_plot.kSIM[line] + 0.003\n",
" if df_plot.method[line] == 'Baseline':\n",
" x_pos = df_plot.kSIM[line] - 0.003\n",
" y_pos = df_plot.kBET[line]\n",
" if df_plot.method[line] == 'iNMF':\n",
" y_pos -= 0.01\n",
" alignment = 'right' if df_plot.method[line] == 'Baseline' else 'left'\n",
" ax.text(x_pos, y_pos, df_plot.method[line], ha = alignment, size = 'medium', color = 'black')\n",
"plt.xlabel('kSIM acceptance rate')\n",
"plt.ylabel('kBET acceptance rate')"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"As this plot shows, a trade-off exists between good mixture of cells (in terms of kBET acceptance rate) and maintaining the biology well (in terms of kSIM acceptance rate). *iNMF* method achieves the best mixture of cells, while its consistency with the ground truth biology is the least. *Harmony*, together with *scVI* and *Scanorama*, have a better balance between the two measurements.\n",
"\n",
"Therefore, in general, the choice of batch correction method really depends on the dataset and your analysis goal."
]
}
],
"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.9.12"
},
"toc": {
"base_numbering": 1,
"nav_menu": {},
"number_sections": true,
"sideBar": true,
"skip_h1_title": false,
"title_cell": "Table of Contents",
"title_sidebar": "Contents",
"toc_cell": false,
"toc_position": {},
"toc_section_display": true,
"toc_window_display": false
}
},
"nbformat": 4,
"nbformat_minor": 4
}