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

# # Python MAGIC EMT tutorial
# 
# ## MAGIC (Markov Affinity-Based Graph Imputation of Cells)
# 
# - MAGIC imputes missing data values on sparse data sets, restoring the structure of the data
# - It also proves dimensionality reduction and gene expression visualizations
# - MAGIC can be performed on a variety of datasets
# - Here, we show the effectiveness of MAGIC on epithelial-to-mesenchymal transition (EMT) data
#     
# Markov Affinity-based Graph Imputation of Cells (MAGIC) is an algorithm for denoising and transcript recover of single cells applied to single-cell RNA sequencing data, as described in Van Dijk D et al. (2018), Recovering Gene Interactions from Single-Cell Data Using Data Diffusion, Cell https://www.cell.com/cell/abstract/S0092-8674(18)30724-4.
# 
# This tutorial shows loading, preprocessing, MAGIC imputation and visualization of myeloid and erythroid cells in mouse bone marrow. You can edit it yourself at https://colab.research.google.com/github/KrishnaswamyLab/MAGIC/blob/master/python/tutorial_notebooks/emt_tutorial.ipynb
# 
# ### Table of Contents
# 
# <a href='#install'>Installation</a>
# <br/>
# <a href='#data'>Loading data</a>
# <br/>
# <a href='#pre'>Data preprocessing</a>
# <br/>
# <a href='#magic'>Running MAGIC</a>
# <br/>
# <a href='#genes'>Visualizing gene-gene interactions</a>
# <br/>
# <a href='#manifold'>Visualizing cell trajectories with PCA on MAGIC</a>
# <br/>
# <a href='#postprocess'>Using MAGIC data in downstream analysis</a>
# 
# <a id="install"></a>
# 
# ### Installation 
# 
# If you haven't yet installed MAGIC, we can install it directly from this Jupyter Notebook.

# In[ ]:


get_ipython().system('pip install --user magic-impute')


# ### Importing MAGIC
# 
# Here, we'll import MAGIC along with other popular packages that will come in handy.

# In[1]:


import magic
import scprep

import numpy as np
import pandas as pd
import matplotlib
import matplotlib.pyplot as plt

# Matplotlib command for Jupyter notebooks only
get_ipython().run_line_magic('matplotlib', 'inline')


# ### Loading Data
# 
# Load your data using one of the following `scprep.io` methods: `load_csv`,`load_tsv`,`load_fcs`,`load_mtx`,`load_10x`. 
# 
# You can read about how to use them with `help(scprep.io.load_csv)` or on https://scprep.readthedocs.io/.

# In[2]:


emt_data = scprep.io.load_csv('../../data/HMLE_TGFb_day_8_10.csv.gz', cell_names=False)
emt_data.head()


# <a id="data"></a>
# 
# ### Data Preprocessing
# 
# The EMT example data we are using is already pre-filtered and nicely distributed, so we will only demonstrate these preprocessing steps, and not actually perform them on the data. However, these steps are essential for performing MAGIC on **single-cell RNA-seq data**, so don't skip them if that's what you're working with. 
# 
# #### Filtering your data
# 
# After loading your data, you're going to want to determine the molecule per cell and molecule per gene cutoffs with which to filter the data, in order to remove lowly expressed genes and cells with a small library size.

# In[3]:


scprep.plot.plot_library_size(emt_data, cutoff=1500)


# In[4]:


# be sure to uncomment this, unless your data is pre-filtered
# emt_data = scprep.filter.filter_library_size(emt_data, cutoff=1500)
emt_data.head()


# #### Normalizing your data
# 
# After filtering, the next steps are to perform library size normalization and transformation. Log transformation is frequently used for single-cell RNA-seq, however, this requires the addition of a pseudocount to avoid infinite values at zero. We instead use a square root transform, which has similar properties to the log transform but has no problem with zeros.

# In[5]:


emt_data = scprep.normalize.library_size_normalize(emt_data)
# be sure to uncomment this, unless your data is pre-transformed
# emt_data = scprep.transform.sqrt(emt_data)
emt_data.head()


# <a id="magic"></a>
# 
# ### Running MAGIC
# 
# Now that your data has been preprocessed, you are ready to run MAGIC.
# 
# #### Creating the MAGIC operator
# 
# If you don't specify parameters, MAGIC creates an operator with the following default values: `knn=5`, `knn_max = 3 * knn`, `decay=1`, `t=3`.

# In[6]:


magic_op = magic.MAGIC()


# #### Running MAGIC with gene selection
# 
# The magic_op.fit_transform function takes the normalized data and an array of selected genes as its arguments. If no genes are provided, MAGIC will return a matrix of all genes. The same can be achieved by substituting the array of gene names with `genes='all_genes'`.

# In[7]:


emt_magic = magic_op.fit_transform(emt_data, genes=['VIM', 'CDH1', 'ZEB1'])


# ### Visualizing gene-gene relationships
# <a id="genes"></a>
# 
# We can see gene-gene relationships much more clearly after applying MAGIC. Note that the change in absolute values of gene expression is not meaningful - the relative difference is all that matters.

# In[8]:


fig, (ax1, ax2) = plt.subplots(1,2, figsize=(16, 6))

scprep.plot.scatter(x=emt_data['VIM'], y=emt_data['CDH1'], c=emt_data['ZEB1'],  ax=ax1,
                    xlabel='VIM', ylabel='CDH1', legend_title="ZEB1", title='Before MAGIC')

scprep.plot.scatter(x=emt_magic['VIM'], y=emt_magic['CDH1'], c=emt_magic['ZEB1'], ax=ax2,
                    xlabel='VIM', ylabel='CDH1', legend_title="ZEB1", title='After MAGIC')

plt.tight_layout()
plt.show()


# The original data suffers from dropout to the point that we cannot infer anything about the gene-gene relationships. As you can see, the gene-gene relationships are much clearer after MAGIC. These relationships also match the biological progression we expect to see in EMT data.

# #### Setting the MAGIC operator parameters
# 
# If you did not specify any parameters for your MAGIC operator, you can do so without going through the hassle of creating a new one using the `magic_op.set_params` method. Since our example EMT dataset is rather small, we can set `knn=3`, rather than the default `knn=5`. 
# 
# 

# In[9]:


magic_op.set_params(knn=3)


# We can now run MAGIC on the data again with the new parameters. Given that we have already fitted our MAGIC operator to the data, we should run the magic_op.transform method.

# In[10]:


emt_magic = magic_op.transform(genes=['VIM', 'CDH1', 'ZEB1'])
emt_magic.head()


# In[11]:


fig, (ax1, ax2) = plt.subplots(1,2, figsize=(16, 6))

scprep.plot.scatter(x=emt_data['VIM'], y=emt_data['CDH1'], c=emt_data['ZEB1'],  ax=ax1,
                    xlabel='VIM', ylabel='CDH1', legend_title="ZEB1", title='Before MAGIC')

scprep.plot.scatter(x=emt_magic['VIM'], y=emt_magic['CDH1'], c=emt_magic['ZEB1'], ax=ax2,
                    xlabel='VIM', ylabel='CDH1', legend_title="ZEB1", title='After MAGIC')

plt.tight_layout()
plt.show()


# <a id="manifold"></a>
# 
# ### Visualizing cell trajectories with PCA on MAGIC

# We can extract the principal components of the smoothed data by passing the keyword `genes='pca_only'` and use this for visualizing the data.

# In[12]:


emt_magic_pca = magic_op.transform(genes='pca_only')
emt_magic_pca.head()


# We'll also run PCA on the raw data to compare.

# In[13]:


from sklearn.decomposition import PCA
emt_pca = PCA(n_components=3).fit_transform(np.array(emt_data))


# In[14]:


fig, (ax1, ax2) = plt.subplots(1,2, figsize=(16, 6))

scprep.plot.scatter2d(emt_pca, c=emt_data['VIM'], 
                      label_prefix="PC", title='PCA without MAGIC',
                      legend_title="VIM", ax=ax1, ticks=False)

scprep.plot.scatter2d(emt_magic_pca, c=emt_magic['VIM'], 
                      label_prefix="PC", title='PCA with MAGIC',
                      legend_title="VIM", ax=ax2, ticks=False)

plt.tight_layout()
plt.show()


# We can also plot this in 3D.

# In[15]:


from mpl_toolkits.mplot3d import Axes3D

fig, (ax1, ax2) = plt.subplots(1,2, figsize=(16, 6), subplot_kw={'projection':'3d'})

scprep.plot.scatter3d(emt_pca, c=emt_data['VIM'], 
                      label_prefix="PC", title='PCA without MAGIC',
                      legend_title="VIM", ax=ax1, ticks=False)

scprep.plot.scatter3d(emt_magic_pca, c=emt_magic['VIM'], 
                      label_prefix="PC", title='PCA with MAGIC',
                      legend_title="VIM", ax=ax2, ticks=False)

plt.tight_layout()
plt.show()


# ### Visualizing MAGIC values with PHATE
# 
# In complex systems, two dimensions of PCA are not sufficient to view the entire space. For this, PHATE is a suitable visualization tool which works hand in hand with MAGIC to view how gene expression evolves along a trajectory. For this, you will need to have installed PHATE. For help using PHATE, visit https://phate.readthedocs.io/.

# In[ ]:


get_ipython().system('pip install --user phate')


# In[16]:


import phate


# In[17]:


data_phate = phate.PHATE().fit_transform(emt_data)


# In[18]:


scprep.plot.scatter2d(data_phate, c=emt_magic['VIM'], figsize=(12,9),
                      ticks=False, label_prefix="PHATE", legend_title="VIM")


# Note that the structure of the data that we see here is slightly more subtle than in PCA and makes more obvious the VIM+ branch. To learn more about PHATE, visit <https://phate.readthedocs.io/>.

# ### Exact vs approximate MAGIC
# 
# If we are imputing many genes at once, we can speed this process up with the argument `solver='approximate'`, which applies denoising in the PCA space and then projects these denoised principal components back onto the genes of interest. Note that this may return some small negative values. You will see below, however, that the results are largely similar to exact MAGIC.

# In[20]:


approx_magic_op = magic.MAGIC(solver="approximate")
approx_emt_magic = approx_magic_op.fit_transform(emt_data, genes='all_genes')


# In[21]:


fig, (ax1, ax2) = plt.subplots(1,2, figsize=(16, 6))

scprep.plot.scatter(x=emt_magic['VIM'], y=emt_magic['CDH1'], c=emt_magic['ZEB1'],  ax=ax1,
                    xlabel='VIM', ylabel='CDH1', legend_title="ZEB1", title='Exact MAGIC')

scprep.plot.scatter(x=approx_emt_magic['VIM'], y=approx_emt_magic['CDH1'], c=approx_emt_magic['ZEB1'], ax=ax2,
                    xlabel='VIM', ylabel='CDH1', legend_title="ZEB1", title='Approximate MAGIC')

plt.tight_layout()
plt.show()


# ### Animating the MAGIC smoothing process

# To visualize what it means to set `t` in MAGIC, we can plot an animation of the smoothing process, from raw to imputed values. Below, we show an animation of _Mpo_, _Klf1_ and _Ifitm1_ with increasingly more smoothing.

# In[20]:


magic.plot.animate_magic(emt_data, gene_x="VIM", gene_y="CDH1", gene_color="ZEB1",
                         operator=magic_op, t_max=12)


# We can also save the animation to a `.gif` file.

# In[22]:


# make the font size bigger
plt.rc('font', size=16)
magic.plot.animate_magic(emt_data, gene_x="VIM", gene_y="CDH1", gene_color="ZEB1",
                         operator=magic_op, t_max=12, filename="magic.gif", dpi=300)


# ### Using the MAGIC smoothed data in downstream analysis
# <a id="postprocess"></a>
# 
# Finally, if you wish to use the full smoothed matrix in any downstream analysis, you can extract it with the keyword `genes="all_genes"`. Note that this matrix may be very large.

# In[23]:


emt_magic = magic_op.transform(genes="all_genes")
emt_magic.head()


# If you wish to export the data, you can do so in many different ways. We recommend the `feather` format, which is fast and lightweight, but `csv` is more commonly used and is compatible with many more software packages.

# In[ ]:


get_ipython().system('pip install --user feather-format')


# In[24]:


import feather
feather.write_dataframe(emt_magic, "emt_magic.feather")


# In[ ]:


emt_magic.to_csv("emt_magic.csv")