#!/usr/bin/env python
# coding: utf-8

# # Preprocessing the data of scRNA-seq with omicverse[GPU]
# 
# The count table, a numeric matrix of genes × cells, is the basic input data structure in the analysis of single-cell RNA-sequencing data. A common preprocessing step is to adjust the counts for variable sampling efficiency and to transform them so that the variance is similar across the dynamic range. 
# 
# Suitable methods to preprocess the scRNA-seq is important. Here, we introduce some preprocessing step to help researchers can perform downstream analysis easyier.
# 
# User can compare our tutorial with [scanpy'tutorial](https://scanpy-tutorials.readthedocs.io/en/latest/pbmc3k.html) to learn how to use omicverse well
# 
# Colab_Reproducibility：https://colab.research.google.com/drive/1DXLSls_ppgJmAaZTUvqazNC_E7EDCxUe?usp=sharing

# ## Installation
# 
# Note that the GPU module is not directly present and needs to be installed separately, for a detailed [tutorial](https://rapids-singlecell.readthedocs.io/en/latest/index.html) see [https://rapids-singlecell.readthedocs.io/en/latest/index.html](https://rapids-singlecell.readthedocs.io/en/latest/index.html)
# 
# ### pip
# ```shell
# pip install rapids-singlecell
# #rapids
# pip install \
#     --extra-index-url=https://pypi.nvidia.com \
#     cudf-cu12==24.4.* dask-cudf-cu12==24.4.* cuml-cu12==24.4.* \
#     cugraph-cu12==24.4.* cuspatial-cu12==24.4.* cuproj-cu12==24.4.* \
#     cuxfilter-cu12==24.4.* cucim-cu12==24.4.* pylibraft-cu12==24.4.* \
#     raft-dask-cu12==24.4.* cuvs-cu12==24.4.*
# #cupy
# pip install cupy-cuda12x
# ```
# 
# ### conda-env
# Note that in order to avoid conflicts, we will consider installing rapid_singlecell first before installing omicverse.
# 
# The easiest way to install rapids-singlecell is to use one of the yaml file provided in the [conda](https://github.com/Starlitnightly/omicverse/tree/master/conda) folder. These yaml files install everything needed to run the example notebooks and get you started.
# ```shell
# conda env create -f conda/omicverse_gpu.yml
# # or
# mamba env create -f conda/omicverse_gpu.yml
# ```

# In[1]:


import omicverse as ov
import scanpy as sc
ov.plot_set()
ov.settings.gpu_init()


# The data consist of 3k PBMCs from a Healthy Donor and are freely available from 10x Genomics ([here](http://cf.10xgenomics.com/samples/cell-exp/1.1.0/pbmc3k/pbmc3k_filtered_gene_bc_matrices.tar.gz) from this [webpage](https://support.10xgenomics.com/single-cell-gene-expression/datasets/1.1.0/pbmc3k)). On a unix system, you can uncomment and run the following to download and unpack the data. The last line creates a directory for writing processed data.

# In[2]:


# !mkdir data
#!wget http://cf.10xgenomics.com/samples/cell-exp/1.1.0/pbmc3k/pbmc3k_filtered_gene_bc_matrices.tar.gz -O data/pbmc3k_filtered_gene_bc_matrices.tar.gz
#!cd data; tar -xzf pbmc3k_filtered_gene_bc_matrices.tar.gz
# !mkdir write


# In[3]:


adata = sc.read_10x_mtx(
    'data/filtered_gene_bc_matrices/hg19/',  # the directory with the `.mtx` file
    var_names='gene_symbols',                # use gene symbols for the variable names (variables-axis index)
    cache=True)                              # write a cache file for faster subsequent reading
adata


# In[4]:


adata.var_names_make_unique()
adata.obs_names_make_unique()


# ## Preprocessing
# 
# ### Quantity control
# 
# For single-cell data, we require quality control prior to analysis, including the removal of cells containing double cells, low-expressing cells, and low-expressing genes. In addition to this, we need to filter based on mitochondrial gene ratios, number of transcripts, number of genes expressed per cell, cellular Complexity, etc. For a detailed description of the different QCs please see the document: https://hbctraining.github.io/scRNA-seq/lessons/04_SC_quality_control.html

# In[5]:


ov.pp.anndata_to_GPU(adata)


# In[6]:


get_ipython().run_cell_magic('time', '', "adata=ov.pp.qc(adata,\n              tresh={'mito_perc': 0.2, 'nUMIs': 500, 'detected_genes': 250},\n              batch_key=None)\nadata\n")


# ### High variable Gene Detection
# 
# Here we try to use Pearson's method to calculate highly variable genes. This is the method that is proposed to be superior to ordinary normalisation. See [Article](https://www.nature.com/articles/s41592-023-01814-1#MOESM3) in *Nature Method* for details.
# 

# normalize|HVGs：We use | to control the preprocessing step, | before for the normalisation step, either `shiftlog` or `pearson`, and | after for the highly variable gene calculation step, either `pearson` or `seurat`. Our default is `shiftlog|pearson`.
# 
# - if you use `mode`=`shiftlog|pearson` you need to set `target_sum=50*1e4`, more people like to se `target_sum=1e4`, we test the result think 50*1e4 will be better
# - if you use `mode`=`pearson|pearson`, you don't need to set `target_sum`
# 
# <div class="admonition warning">
#   <p class="admonition-title">Note</p>
#   <p>
#     if the version of `omicverse` lower than `1.4.13`, the mode can only be set between `scanpy` and `pearson`.
#   </p>
# </div>
# 

# In[7]:


get_ipython().run_cell_magic('time', '', "adata=ov.pp.preprocess(adata,mode='shiftlog|pearson',n_HVGs=2000,)\nadata\n")


# Set the .raw attribute of the AnnData object to the normalized and logarithmized raw gene expression for later use in differential testing and visualizations of gene expression. This simply freezes the state of the AnnData object.

# In[8]:


adata.raw = adata
adata = adata[:, adata.var.highly_variable_features]
adata


# ## Principal component analysis
# 
# In contrast to scanpy, we do not directly scale the variance of the original expression matrix, but store the results of the variance scaling in the layer, due to the fact that scale may cause changes in the data distribution, and we have not found scale to be meaningful in any scenario other than a principal component analysis

# In[9]:


get_ipython().run_cell_magic('time', '', 'ov.pp.scale(adata)\nadata\n')


# If you want to perform pca in normlog layer, you can set `layer`=`normlog`, but we think scaled is necessary in PCA.

# In[10]:


get_ipython().run_cell_magic('time', '', "ov.pp.pca(adata,layer='scaled',n_pcs=50)\nadata\n")


# In[11]:


adata.obsm['X_pca']=adata.obsm['scaled|original|X_pca']
ov.utils.embedding(adata,
                  basis='X_pca',
                  color='CST3',
                  frameon='small')


# ## Embedding the neighborhood graph
# 
# We suggest embedding the graph in two dimensions using UMAP (McInnes et al., 2018), see below. It is potentially more faithful to the global connectivity of the manifold than tSNE, i.e., it better preserves trajectories. In some ocassions, you might still observe disconnected clusters and similar connectivity violations. They can usually be remedied by running:

# In[23]:


get_ipython().run_cell_magic('time', '', "ov.pp.neighbors(adata, n_neighbors=15, n_pcs=50,\n               use_rep='scaled|original|X_pca',method='cagra')\n")


# To visualize the PCA’s embeddings, we use the `pymde` package wrapper in omicverse. This is an alternative to UMAP that is GPU-accelerated.

# In[19]:


adata.obsm["X_mde"] = ov.utils.mde(adata.obsm["scaled|original|X_pca"])
adata


# In[20]:


ov.pl.embedding(adata,
                basis='X_mde',
                color='CST3',
                frameon='small')


# You also can use `umap` to visualize the neighborhood graph

# In[21]:


ov.pp.umap(adata)


# In[22]:


ov.pl.embedding(adata,
                basis='X_umap',
                color='CST3',
                frameon='small')


# ## Clustering the neighborhood graph
# 
# As with Seurat and many other frameworks, we recommend the Leiden graph-clustering method (community detection based on optimizing modularity) by Traag *et al.* (2018). Note that Leiden clustering directly clusters the neighborhood graph of cells, which we already computed in the previous section.

# In[24]:


ov.pp.leiden(adata)


# In[30]:


ov.pp.anndata_to_CPU(adata)


# We redesigned the visualisation of embedding to distinguish it from scanpy's embedding by adding the parameter `fraemon='small'`, which causes the axes to be scaled with the colourbar

# In[25]:


ov.pl.embedding(adata,
                basis='X_mde',
                color=['leiden', 'CST3', 'NKG7'],
                frameon='small')


# We also provide a boundary visualisation function `ov.utils.plot_ConvexHull` to visualise specific clusters.
# 
# Arguments: 
# - color: if None will use the color of clusters
# - alpha: default is 0.2

# In[26]:


import matplotlib.pyplot as plt
fig,ax=plt.subplots( figsize = (4,4))

ov.pl.embedding(adata,
                basis='X_mde',
                color=['leiden'],
                frameon='small',
                show=False,
                ax=ax)

ov.pl.ConvexHull(adata,
                basis='X_mde',
                cluster_key='leiden',
                hull_cluster='0',
                ax=ax)


# If you have too many labels, e.g. too many cell types, and you are concerned about cell overlap, then consider trying the `ov.utils.gen_mpl_labels` function, which improves text overlap.
# In addition, we make use of the `patheffects` function, which makes our text have outlines
# 
# - adjust_kwargs: it could be found in package `adjusttext`
# - text_kwargs: it could be found in class `plt.texts`

# In[27]:


from matplotlib import patheffects
import matplotlib.pyplot as plt
fig, ax = plt.subplots(figsize=(4,4))

ov.pl.embedding(adata,
                  basis='X_mde',
                  color=['leiden'],
                   show=False, legend_loc=None, add_outline=False, 
                   frameon='small',legend_fontoutline=2,ax=ax
                 )

ov.utils.gen_mpl_labels(
    adata,
    'leiden',
    exclude=("None",),  
    basis='X_mde',
    ax=ax,
    adjust_kwargs=dict(arrowprops=dict(arrowstyle='-', color='black')),
    text_kwargs=dict(fontsize= 12 ,weight='bold',
                     path_effects=[patheffects.withStroke(linewidth=2, foreground='w')] ),
)


# In[28]:


marker_genes = ['IL7R', 'CD79A', 'MS4A1', 'CD8A', 'CD8B', 'LYZ', 'CD14',
                'LGALS3', 'S100A8', 'GNLY', 'NKG7', 'KLRB1',
                'FCGR3A', 'MS4A7', 'FCER1A', 'CST3', 'PPBP']


# In[29]:


sc.pl.dotplot(adata, marker_genes, groupby='leiden',
             standard_scale='var');


# ## Finding marker genes
# 
# Let us compute a ranking for the highly differential genes in each cluster. For this, by default, the .raw attribute of AnnData is used in case it has been initialized before. The simplest and fastest method to do so is the t-test.

# In[31]:


sc.tl.dendrogram(adata,'leiden',use_rep='scaled|original|X_pca')
sc.tl.rank_genes_groups(adata, 'leiden', use_rep='scaled|original|X_pca',
                        method='t-test',use_raw=False,key_added='leiden_ttest')
sc.pl.rank_genes_groups_dotplot(adata,groupby='leiden',
                                cmap='Spectral_r',key='leiden_ttest',
                                standard_scale='var',n_genes=3)


# cosg is also considered to be a better algorithm for finding marker genes. Here, omicverse provides the calculation of cosg
# 
# Paper: [Accurate and fast cell marker gene identification with COSG](https://academic.oup.com/bib/advance-article-abstract/doi/10.1093/bib/bbab579/6511197?redirectedFrom=fulltext)
# 
# Code: https://github.com/genecell/COSG
# 

# In[32]:


sc.tl.rank_genes_groups(adata, groupby='leiden', 
                        method='t-test',use_rep='scaled|original|X_pca',)
ov.single.cosg(adata, key_added='leiden_cosg', groupby='leiden')
sc.pl.rank_genes_groups_dotplot(adata,groupby='leiden',
                                cmap='Spectral_r',key='leiden_cosg',
                                standard_scale='var',n_genes=3)


# ## Other plotting
# 
# Next, let's try another chart, which we call the Stacked Volcano Chart. We need to prepare two dictionaries, a `data_dict` and a `color_dict`, both of which have the same key requirements.
# 
# For `data_dict`. we require the contents within each key to be a DataFrame containing ['names','logfoldchanges','pvals_adj'], where names stands for gene names, logfoldchanges stands for differential expression multiplicity, pvals_adj stands for significance p-value
# 

# In[33]:


data_dict={}
for i in adata.obs['leiden'].cat.categories:
    data_dict[i]=sc.get.rank_genes_groups_df(adata, group=i, key='leiden_ttest',
                                            pval_cutoff=None,log2fc_min=None)


# In[34]:


data_dict.keys()


# In[35]:


data_dict[i].head()


# For `color_dict`, we require that the colour to be displayed for the current key is stored within each key.`

# In[36]:


type_color_dict=dict(zip(adata.obs['leiden'].cat.categories,
                         adata.uns['leiden_colors']))
type_color_dict


# There are a number of parameters available here for us to customise the settings. Note that when drawing stacking_vol with omicverse version less than 1.4.13, there is a bug that the vertical coordinate is constant at [-15,15], so we have added some code in this tutorial for visualisation.
# 
# - data_dict: dict, in each key, there is a dataframe with columns of ['logfoldchanges','pvals_adj','names']
# - color_dict: dict, in each key, there is a color for each omic
# - pval_threshold: float, pvalue threshold for significant genes
# - log2fc_threshold: float, log2fc threshold for significant genes
# - figsize: tuple, figure size
# - sig_color: str, color for significant genes
# - normal_color: str, color for non-significant genes
# - plot_genes_num: int, number of genes to plot
# - plot_genes_fontsize: int, fontsize for gene names
# - plot_genes_weight: str, weight for gene names

# In[37]:


fig,axes=ov.utils.stacking_vol(data_dict,type_color_dict,
            pval_threshold=0.01,
            log2fc_threshold=2,
            figsize=(8,4),
            sig_color='#a51616',
            normal_color='#c7c7c7',
            plot_genes_num=2,
            plot_genes_fontsize=6,
            plot_genes_weight='bold',
            )

#The following code will be removed in future
y_min,y_max=0,0
for i in data_dict.keys():
    y_min=min(y_min,data_dict[i]['logfoldchanges'].min())
    y_max=max(y_max,data_dict[i]['logfoldchanges'].max())
for i in adata.obs['leiden'].cat.categories:
    axes[i].set_ylim(y_min,y_max)
plt.suptitle('Stacking_vol',fontsize=12)   


# In[ ]: