# ---
# jupyter:
#   jupytext:
#     text_representation:
#       extension: .py
#       format_name: light
#       format_version: '1.5'
#       jupytext_version: 1.14.4
#   kernelspec:
#     display_name: Python 3
#     language: python
#     name: python3
# ---

# # SCALE tutorial

# ### Read alignment and peak calling
#
# ** We used Kundaje pipeline (https://github.com/kundajelab/atac_dnase_pipelines) for aligning pair ended reads to hg19 and removing duplicates as scABC (https://github.com/SUwonglab/scABC/) did **  
# bds_scr [SCR_NAME] atac.bds -align -species hg19 -species_file [SPECIES_FILE_PATH] -nth [NUM_THREADS] -fastq1_1 [READ_PAIR1] -fastq1_2 [READ_PAIR2]
#
# ** The resulting bam files were merged using samtools. **  
# samtools merge [AGGREGATE_BAM] *.trim.PE2SE.nodup.bam  
#
# ** The merged bam file was then used as input into MACS2 to call merged peaks for later analysis, using the Kundaje pipeline again.**  
# bds_scr [SCR_NAME] atac.bds -species hg19 -species_file [SPECIES_FILE_PATH] -nth [NUM_THREADS] -se -filt_bam [AGGREGATE_BAM]
#

# ### Preprocessing
# from scale.utils import sample_filter, peak_filter, cell_filter  
# ** We filtered peaks that presented in >= 10 cells with >= 2 reads by peak_filter function **  
# data = peak_filter(data)    
# ** We filtered cells that >= number of peaks / 50 by cell_filter function **  
# data = cell_filter(data)  
# ** Or combine peak_filter and cell_filter **  
# data = sample_filter(data)

# # Get Started from a scATAC-seq data matrix

# ## Running SCALE model 
# e.g. SCALE -d data_matrix -k cluster_number -o output_dir  
# ** input:**  
# peak count matrix
#
# ** output: **  
# 1. feature.txt 
# 2. cluster_assignments.txt
# 3. impute_data.txt
# 4. model.pt

# +
# # !SCALE.py -d data.txt -k 6 -e 1000 -o output --lr 0.002
# -

# ## Analysis workflow
# 1. visualization of latent feature embeddings
# 2. confusion matrix
# 3. data imputation analysis
# 4. cell type specific elements

# +
import pandas as pd
import numpy as np
from matplotlib import pyplot as plt
import seaborn as sns
from sklearn.metrics import confusion_matrix

from scale.plot import plot_confusion_matrix, plot_embedding, plot_heatmap
from scale.utils import read_labels, reassign_cluster_with_ref, pairwise_pearson, sort_by_mad, peak_selection, get_decoder_weight
from scale.specifity import *

# %load_ext autoreload
# %autoreload 2
# %matplotlib inline

# +
# plt.rcParams['savefig.dpi'] = 300
# plt.rcParams['figure.dpi'] = 300
# -

input_dir = '../data/'
out_dir = '../output/'

# +
feature_file = out_dir+'feature.txt'
assignments_file = out_dir+'cluster_assignments.txt'
impute_data_file = out_dir+'imputed_data.txt'

raw_data_file = input_dir + 'data.txt'
reference_file = input_dir + 'labels.txt'
# -

raw_data = pd.read_csv(raw_data_file, sep='\t', index_col=0)
imputed_data = pd.read_csv(impute_data_file, sep='\t', index_col=0)

# +
ref, classes, le = read_labels(reference_file, return_enc=True)
classes = ['Mono', 'LMPP', 'LSC_SU353', 'LSC_SU070', 'Blast_SU353', 'Blast_SU070']
classes_ = le.transform(classes) # class number name

pred, pred_classes = read_labels(assignments_file)
pred = reassign_cluster_with_ref(pred, ref)

y = le.inverse_transform(ref)
pred_ = le.inverse_transform(pred)
y_pred = np.array(['cluster'+str(c+1) for c in pred])
# _, assignments = cluster_acc(assignments, ref) # reassignment cluster
# -

# ## 1. Feature embedding
# different color representing ground truth cell types

feature = pd.read_csv(feature_file, sep='\t', index_col=0, header=None)

legend_params = {'loc':'right', 'ncol':1, 'bbox_to_anchor':(0.4, 0.2), 'fontsize':8, 'markerscale':1.5}
plot_embedding(feature, y, classes=classes, markersize=10, figsize=(4,4), legend_params=legend_params )

plot_heatmap(feature.T, y, classes=classes, 
             figsize=(8, 3), cmap='RdBu_r', #vmax=8, vmin=-8,
             ylabel='Feature dimension', yticklabels=np.arange(10)+1, 
             cax_title='Feature value',
             row_cluster=False, legend_font=6, 
             center=0)

weight = get_decoder_weight(out_dir+'model.pt')
weight_index = sort_by_mad(raw_data.T).columns
specific_peaks = peak_selection(weight, weight_index)

for i, peak_index in enumerate(specific_peaks): # show the represented peaks of last two components of feature
    if i in [8, 9]:
        for data in [raw_data, imputed_data]:
            peak_data = data.loc[peak_index]
            plot_heatmap(peak_data, y, classes=classes,
                     cmap='Reds', 
                     figsize=(8,3), 
                     cax_title='Peak value', 
                     ylabel='{} peaks of feature {}'.format(len(peak_index), i+1),
                     vmax=1, vmin=0, legend_font=8,
                     row_cluster=False,
                     show_legend=True,
                     show_cax = True,
                     bbox_to_anchor=(0.4, 1.32),
                    )

# ## 2. Confusion matrix 
# cluster assignments compared with ground truth

cm = confusion_matrix(pred_, y, labels=classes)
plot_confusion_matrix(cm, classes, pred_classes+1, 
                    title='', normalize=False, figsize=(3,3), show_cbar=False)

# ## 3. Imputation

# #### Cells correlation of imputed data with meta-cells of each cell type

cell_corr = []
for c in classes:
    cells = np.where(y==c)[0]
    cell_corr.append(pairwise_pearson(imputed_data.iloc[:, cells].T, raw_data.iloc[:, cells].sum(1)))

g = sns.boxplot(data=cell_corr, width=0.5)
sns.despine(offset=10, trim=True)
g.set_xticklabels(classes, rotation=45)
plt.show()

# ## 4. Cell type specific peaks

score_mat = mat_specificity_score(imputed_data, pred)

peak_index, peak_labels = cluster_specific(score_mat, classes=classes_, top=200);print(len(peak_index))
row_labels = ['cluster'+str(c+1) for c in peak_labels]

plot_heatmap(imputed_data.iloc[peak_index], y=y, classes=classes, y_pred=y_pred, row_labels=row_labels, 
             ncol=3,cmap='Reds', vmax=1, row_cluster=False, legend_font=6, cax_title='Peak Value',
             figsize=(6, 8), bbox_to_anchor=(0.4, 1.2), position=(0.8, 0.76, 0.1, 0.015))



StarcoderdataJupyterScriptsDedupFiltered