# ---
# 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
# ---

# +
from pathlib import Path
from collections import Counter

import anndata
import scanpy as sc
import pandas as pd
import seaborn as sns
import umap

import pollock
from pollock.models.model import predict_from_anndata, embed_from_anndata
from pollock.models.explain import explain_predictions
# -

# Specify which module to use for cell type predictions. If you haven't trained a module yet try building one with the [pbmc module training](https://github.com/ding-lab/pollock/blob/master/examples/pbmc_model_training.ipynb) example

module_filepath = 'modules/pbmc'

# download adata object here
#
# it is also the same anndata object create here

adata = sc.read_h5ad('data/pbmc.h5ad')
adata

# create a new anndata object where raw counts are stored in .X

prediction_adata = sc.read_10x_mtx('data/filtered_gene_bc_matrices/hg19/', var_names='gene_symbols')
prediction_adata.var_names_make_unique()
prediction_adata

# ## Module prediction

# predict cell types

predictions = predict_from_anndata(prediction_adata.copy(), module_filepath)
predictions

# merge the predictions with our pbmc anndata object
adata.obs = pd.merge(adata.obs, predictions, left_index=True, right_index=True)
adata.obs

adata.obs['annotated_cell_type'] = adata.obs['leiden'].to_list()

sc.pl.umap(adata, color=['annotated_cell_type', 'predicted_cell_type', 'cell_type_probability'],
          frameon=False, ncols=1)

# ## visualization of cell embeddings

# with pollock we can also visualize the cell embeddings created by the BVAE

cell_embeddings = embed_from_anndata(prediction_adata.copy(), module_filepath)
cell_embeddings

adata.obsm['cell_embeddings'] = cell_embeddings.loc[adata.obs.index].values

# replace pca generated nearest neighbor and umap with cell embedding generated one

adata.obsm['X_umap'] = umap.UMAP().fit_transform(adata.obsm['cell_embeddings'])

sc.pl.umap(adata, color=['annotated_cell_type', 'predicted_cell_type', 'cell_type_probability'],
          frameon=False, ncols=1)

# ## Module explain

# For model explaination, we need to split the dataset up into two pieces
#   - explain : anndata object containing cells to be explained. the explaination algorithm can be time consuming (~10min for 100 cells).
#   - background : anndata object containing background cells, generally this can be any group of cells from the dataset, as long as the are **not** in the explain_adata.

# +
# here we grab ten of each cell type for explain, and put the rest into background
n = 10
cell_types = sorted(set(adata.obs['predicted_cell_type']))
cell_ids = [cell_id for cell_type in cell_types
           for cell_id in adata.obs[adata.obs['predicted_cell_type']==cell_type].index[:n]]

explain = prediction_adata[cell_ids]
explain.obs = adata.obs.loc[cell_ids]

background_mask = [True if cell_id not in explain.obs.index else False
                  for cell_id in prediction_adata.obs.index]
background = prediction_adata[background_mask]

Counter(explain.obs['predicted_cell_type']).most_common()
# -

# There are some additional arguments that need be be given to explain_predictions
# - module_filepath : filepath of module used for prediction.
# - prediction_key : column in explain.obs where cell type predictions are stored.
# - n_background_cells : number of cells to sample from background_adata. A larger number means more accurate results, but a longer runtime. The default of 100 is usually sufficient.
#
# Gene weights are returned from explain_predictions. They are represented as a dataframe where columns are genes, rows are cells, and the values correspond to how much each gene contributed to a given cells classification. High values mean that gene pushed that cell towards a given prediction.

weights = explain_predictions(explain.copy(), background.copy(), module_filepath,
                             prediction_key='predicted_cell_type', n_background_cells=100)

weights

# visualizing known markers for pbmc cell types

# constructing weights as an AnnData object so we can use anndata/scanpy plotting functionality
feature_adata = anndata.AnnData(X=weights.values,
                       obs=explain.obs.loc[weights.index])
feature_adata.var.index = weights.columns
feature_adata

markers = {
    'NK': ['GNLY', 'NKG7'],
    'T-cell': ['CD3D', 'CD8A'],
    'B-cell': ['CD79A', 'MS4A1'],
    'Monocytes': ['FCGR3A', 'CD14', 'ITGAM'],
    'Megakaryocytes': ['ITGA2B']
}

sc.pl.heatmap(feature_adata, var_names=markers, groupby='predicted_cell_type')

# Here we run scanpy's DEG workflow to see what features are differentially weighted among the different cell types

adata = feature_adata.copy()
sc.pp.scale(adata)
adata

sc.tl.rank_genes_groups(adata, groupby='predicted_cell_type', method='wilcoxon')
sc.pl.rank_genes_groups_heatmap(adata)

# Viewing specific groups of cells

sc.pl.rank_genes_groups_dotplot(adata, groups=['B', 'Megakaryocytes'])



StarcoderdataJupyterScriptsDedupFiltered