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

# In[ ]:


# get the latest version from pypi 
# for other installations approaches, see https://dynamo-release.readthedocs.io/en/latest/ten_minutes_to_dynamo.html#how-to-install
get_ipython().system('pip install dynamo-release --upgrade --quiet')


# Let us show that velocity estimation based on slam-seq data gives more consistent and clean results. 
# Note that this is irrespective of the velocity tools you will use. Original scSLAM-seq authors show that velocyto gives better results when using slam-seq data. Here I am using the same result with the recently improved scvelo package. Please come back to see the result with dynamo! 

# In[1]:


import pandas as pd 
import numpy as np 

import dynamo as dyn
dyn.get_all_dependencies_version()


# In[2]:


# emulate ggplot2 plotting styple with white background

dyn.configuration.set_figure_params('dynamo', background='white')  


# # data from Hendriks et al. 2018 

# Let us first run velocity estimation based on splicing data
# 
# use wget to download the raw_data_loom_combined file and then load the loom data 
# 
# wget https://www.dropbox.com/s/a9ozcynpxudqdis/raw_data_loom_combined.loom?dl=1
# 
# here I just use data I have in my local directory

# In[3]:


adata_MCMV = dyn.read_loom('/Users/xqiu/Dropbox (Personal)/dynamo/notebook_data/raw_data_loom_combined.loom')


# In[4]:


# add a treatment label to the adata_MCMV object
import re
batch = list()

for cell in adata_MCMV.obs.index.values:
    regex = re.compile('3LF70')
    result = regex.search(cell)
    if result is not None:
        label = 'mock'
    else:
        label = 'mcmv'
    batch.append(label)

# add it to the adata_MCMV object
adata_MCMV.obs['virus'] = np.array(batch)


# In[5]:


adata_MCMV


# In[6]:


dyn.pp.recipe_monocle(adata_MCMV, n_top_genes=1000)
dyn.tl.dynamics(adata_MCMV, model='stochastic') 
# or dyn.tl.dynamics(adata, model='stochastic')
# or dyn.tl.dynamics(adata, model='stochastic', est_method='negbin')

dyn.tl.reduceDimension(adata_MCMV)

dyn.pl.umap(adata_MCMV, color='virus')


# In[7]:


dyn.tl.cell_velocities(adata_MCMV)
# dyn.pl.phase_portraits(adata_MCMV, genes=['RERE', 'ENO1', 'DHRS3'], ncols=3, figsize=(3, 3))
dyn.pl.cell_wise_vectors(adata_MCMV, color=['virus'], basis='umap', quiver_scale=2, show_legend='on data') # ['GRIA3', 'LINC00982', 'AFF2']


# In[8]:


dyn.pl.grid_vectors(adata_MCMV, color=['virus'], basis='umap', xy_grid_nums=(60,60)) # ['GRIA3', 'LINC00982', 'AFF2']


# In[9]:


dyn.pl.streamline_plot(adata_MCMV, color=['virus'], basis='umap', method='gaussian', density=2) # ['GRIA3', 'LINC00982', 'AFF2']


# In[10]:


# dyn.pl.phase_portraits(adata_MCMV, genes=['Npc2', 'Psme2b', 'Fosl2', 'Mcl1'], ncols=3, figsize=(3, 3))
dyn.vf.VectorField(adata_MCMV, basis='umap')


# In[11]:


dyn.pl.topography(adata_MCMV, color=['virus'], basis='umap', terms=('streamline', 'fixed_points', 'trajectory'), init_state = np.array([10, 0]))


# The above velocity estimation results are kind messy. Next, let us try velocity estimation based on scSLAM-seq data. 
# 
# Note that we use total RNA (labelled new RNA + unlabeled old RNA) as **spliced** while the new RNA as **unspliced**

# In[12]:


tot_RNA = pd.read_csv('https://www.dropbox.com/s/skgesrran9d48oy/emat_tot.txt?dl=1', index_col=0, delimiter='\s')
new_RNA = pd.read_csv('https://www.dropbox.com/s/kz0xj8hw4tbab9r/smat_new.txt?dl=1', index_col=0, delimiter='\s')

from anndata import AnnData
adata_sc_slamseq_MCMV = AnnData(tot_RNA.values.T,
    layers=dict(
        unspliced=new_RNA.values.T,
        spliced = tot_RNA.values.T)) 
adata_sc_slamseq_MCMV

# again, let annotate cells by virus infection
batch = list()

for cell in tot_RNA.columns.values:
    regex = re.compile('mock')
    result = regex.search(cell)
    if result is not None:
        label = 'mock'
    else:
        label = 'mcmv'
    batch.append(label)
    
# add it to the vlm object
adata_sc_slamseq_MCMV.obs['virus'] = np.array(batch)
adata_sc_slamseq_MCMV


# In[13]:


dyn.pp.recipe_monocle(adata_sc_slamseq_MCMV, n_top_genes=583, normalized=True, fg_kwargs={'shared_count': 0})
adata_sc_slamseq_MCMV


# In[14]:


adata_sc_slamseq_MCMV.layers['X_unspliced'], adata_sc_slamseq_MCMV.layers['X_spliced'] = adata_sc_slamseq_MCMV.layers['unspliced'], adata_sc_slamseq_MCMV.layers['spliced']


# In[15]:


dyn.tl.dynamics(adata_sc_slamseq_MCMV, model='deterministic') 
dyn.tl.reduceDimension(adata_sc_slamseq_MCMV)

dyn.pl.umap(adata_sc_slamseq_MCMV, color='virus')


# In[16]:


dyn.pl.cell_wise_vectors(adata_sc_slamseq_MCMV, color=['virus'], basis='umap', quiver_size=2, show_legend='on data') # ['GRIA3', 'LINC00982', 'AFF2']


# In[17]:


dyn.pl.grid_vectors(adata_sc_slamseq_MCMV, color=['virus'], basis='umap', method='gaussian') # ['GRIA3', 'LINC00982', 'AFF2']


# In[18]:


dyn.pl.streamline_plot(adata_sc_slamseq_MCMV, color=['virus'], basis='umap', method='gaussian') # ['GRIA3', 'LINC00982', 'AFF2']


# In[19]:


dyn.vf.VectorField(adata_sc_slamseq_MCMV, basis='umap')


# In[20]:


dyn.pl.topography(adata_sc_slamseq_MCMV, color=['virus'], basis='umap', terms=('streamline', 'fixed_points', 'trajectory'))


# In[21]:


dyn.pl.cell_wise_vectors(adata_sc_slamseq_MCMV, color=['virus'], basis='umap', quiver_size=4, quiver_length=5, show_legend='on data') # ['GRIA3', 'LINC00982', 'AFF2']


# In[22]:


dyn.pl.cell_wise_vectors(adata_sc_slamseq_MCMV, color=['virus'], basis='umap', 
                          quiver_size=4, quiver_length=5, show_legend='on data', background='black') # ['GRIA3', 'LINC00982', 'AFF2']


# In[23]:


dyn.pl.streamline_plot(adata_sc_slamseq_MCMV, color=['virus'], basis='umap', show_legend='on data', background='black') # ['GRIA3', 'LINC00982', 'AFF2']


# the velocity estimation here is much more apparent and consistent. let us save the result