##############
### README ###
##############
# Datafile S2  : Python script used in this study

## Description of the scripts
### _____  
- Script used to perform the trajectory modeling analysis based on a set of signature genes using the STREAM algorithm.
- Input files include the
	- cells’ expression data and metadata files.
	- list of signature genes on which the modeling should be based.
	- list of metabolism genes. 
- Output files include
	- plots showing the trajectory modeled on which the clusters and the motility score were highlighted (figures presented in Figure 2b).
	- list of metabolism genes whose expression varies between adjacent branches of the trajectory. 


##################################
### Script for STREAM analysis ###
##################################

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

## STREAM ANALYSIS
# Construction of a trajectory based on the motility signature using STREAM.


from datetime import date 
import stream as st
import pandas as pd
import numpy as np
import matplotlib.pyplot as plt
import scipy
import anndata as ad
import statsmodels.api as sm
from rpy2.robjects.packages import importr
from rpy2.robjects import pandas2ri
from rpy2 import robjects
from pylab import *

today = date.today()

st.set_figure_params(dpi=80,style='white',figsize=[5.4,4.8],
                     rc={'image.cmap': 'YlOrRd'})


## Import data

adata=st.read(file_name='../Galaxy7-[ExpressionMatrix.tsv].tabular', 
              file_format="tsv", workdir='./stream_result_' + str(today))
#add metadata
st.add_metadata(adata,
                file_name = "../Galaxy6-[Metadata.csv]_forStream.tsv",
                delimiter='\t')

## Calculate QC - Filtering low-quality genes

st.cal_qc(adata,assay='rna')
st.filter_features(adata,min_n_cells = 5)


## Variable features

st.select_variable_genes(adata,loess_frac=0.04,percentile=95)
adata.uns['var_genes_selected'] = adata.uns['var_genes']


# Consider signature genes as variable genes in order to perform dimensional reduction and trajectories only based on these genes.

adata.uns['var_genes'] = pandas.core.indexes.base.Index(['ACTN4', 'TNC', 'THBS1', 'TGFB1', 'PXN', 'SMAD3', 'SPARCL1', 'TLN1', 'VCL', 'PTK2'])
var_genes = result = [idx for idx, val in enumerate(adata.var_names) if val in adata.uns['var_genes']]
adata.obsm['var_genes'] = adata.X[:,var_genes].copy()

## Dimensional reduction

st.dimension_reduction(adata, method = 'se', feature = 'var_genes', n_neighbors = 50, n_components = 4)

st.plot_dimension_reduction(adata, color=['Cluster','MotilityScore'], n_components = 3, show_graph = False, show_text = False)


st.plot_visualization_2D(adata, method = 'umap', n_neighbors = 50,
                         color = ['MotilityScore', 'label'], use_precomputed = False)


## Trajectory inference


#by setting use_vis=True, we use the manifold from `plot_visualization_2D()` to infer trajectories
st.seed_elastic_principal_graph(adata,n_clusters=10,use_vis=True)
st.plot_dimension_reduction(adata,color=['TumScore_5TF','label','MigrationScore'],n_components=2,show_graph=True,show_text=False)
st.plot_branches(adata,show_text=True)


st.plot_flat_tree(adata,color=['label','MotilityScore','branch_id_alias', 'S4_pseudotime', 'S5_pseudotime'],
                  dist_scale=0.5,show_graph=True,show_text=True)

st.plot_stream_sc(adata,root='S5',color=['label','MotilityScore', 'branch_id_alias'],
                  dist_scale=0.5,show_graph=True,show_text=True)
st.plot_stream(adata, color=['label','MotilityScore'], root = 'S5')


# Import metabolism gene list and add it to the python object

metaGenes = pandas.read_csv("../KEGG2019-Galaxy5-[metabolismGenes].csv", sep=";")
adata.uns['metabolismGenes'] = pandas.core.indexes.base.Index(metaGenes[metaGenes.columns[1]])
adata.uns['metabolismGenes'] = adata.uns['metabolismGenes'][adata.uns['metabolismGenes'].isin(adata.var_names)]
adata.uns['metabolismGenes']


## Diverging genes
# i.e. genes important in defining branching points that are differentially expressed between diverging branches.
# Users first need to choose a pair of adjacent branches, between which differentially expressed genes were calculated by STREAM. Detect differentially expressed markers between different sub-branches.

st.detect_de_markers(adata, marker_list = adata.uns['metabolismGenes'], cutoff_zscore = 1, cutoff_logfc = 0.15,
                     root = 'S5', n_jobs = 2, use_precomputed = False)

st.plot_de_markers(adata)