util.py

import itertools
import os
import pickle
from pathlib import Path
from typing import Dict, List, Union

import matplotlib.pyplot as plt
import numpy as np
import pandas as pd
from scanpy import AnnData
from scanpy.plotting import umap
from scipy.stats import median_abs_deviation
import scanpy as sc


def print_value_counts(df, filehandle=None):
    writer = print if filehandle is None else filehandle.write
    for col in df.columns:
        if 'cell' in col.lower():
            continue
        if df[col].dtype in ['object', 'category']:
            writer(df[col].value_counts().to_string())
            writer('\n----------------------\n')


def summarize_samples(df, sample_name, output_file):
    # Group the DataFrame by the `sample_name` column
    grouped = df.groupby(sample_name)

    summary_dict = {}

    # Iterate through the groups and extract sample names
    for category, group_df in grouped:
        sample_names = ', '.join(group_df['sample'].unique())
        summary_dict[category] = sample_names

    # Create a DataFrame from the summarized results
    summary_df = pd.DataFrame(list(summary_dict.items()), columns=['tumor/normal', 'sample_names'])

    # Export the summarized DataFrame to a CSV file
    summary_df.to_csv(output_file, index=False)


# QC
def is_outlier(adata, metric: str, nmads: int):
    M = adata.obs[metric]
    outlier = (M < np.median(M) - nmads * median_abs_deviation(M)) | (
            np.median(M) + nmads * median_abs_deviation(M) < M
    )
    return outlier


def read_marker_from_file(filename: str, force_update=False, path='markers') -> dict:
    """
    The file should be dataframe with at least two columns: cell_name and Symbol

    Examples\n
    |cell_name|PMID|Symbol|\n
    |Acinar cell|31221802|"PPY, SOD2, GCG, PDK4, SST" (allow duplicate)\n
    |Acinar cell|31273297|"CTRB1, PRSS1, REG1B, CTRB2"\n
    |Acinar cell|32094658|"PRSS1"\n
    """
    cache_file = Path(f'cache/{filename.split("/")[-1].split(".")[0]}.pkl')  # hard coded
    if os.path.exists(cache_file) and not force_update:
        with open(cache_file, 'rb') as f:
            print(f'Load marker gene from {cache_file}')
            return pickle.load(f)
    path = Path(os.path.dirname(__file__)) / Path(path)
    markers = pd.read_csv(path / filename)
    print('Read marker genes from file ...')
    print(markers.head())
    marker_gene_dict = {}
    markers['cell_name'] = markers['cell_name'].str.capitalize()

    for cell_name in markers['cell_name'].unique():
        symbols = markers.loc[markers['cell_name'] == cell_name, 'Symbol'].str.split(', ').explode().tolist()
        # remove duplicate
        marker_gene_dict[cell_name] = list(set(symbols))

    try:
        os.makedirs('cache')
    except FileExistsError:
        pass
    with open(cache_file, 'wb') as f:
        pickle.dump(marker_gene_dict, f)

    return marker_gene_dict


# gene selection
def get_marker_gene_dict() -> dict:
    marker_gene_dict = {
        "Mast": ["TPSAB1", "CPA3", "TPSD1", "MS4A2", "KIT", "SIGLEC6", "SIGLEC8",
                 "CD22", "CSF1", "TIMP3", "GATA2", "ACSL4", "SLC18A2", "CTSG"],
        "pDC": ["LILRA4", "SLC15A4", "PLD4", "CCDC50", "IL3RA", "LY9", "SELL", "GAS6"],
        "cDC1": ["BATF3", "XCR1", "CLEC9A", "PTDSS1", "SCARB1", "IL6ST",
                 "CD40", "TNFRSF10B", "IDO1", "CST7", "CLIC2", "NET1", "ANXA6"],
        "cDC2": ["CD1C", "CLEC10A", "FCGR2B", "ADAM8", "FCER1A", "AXL",
                 "ADAM28", "LY86", "TIMM13", "ARAF"],
        "CD14+ Mono": ["FCN1", "CD14", "CD36", "SELL", "S100A8", "S100A12",
                       "CLEC12A", "MS4A6A", "CXCL14"],
        "CD16+ Mono": ["FCN1", "FCGR3A", "FCGR3B", "LILRA/B1", "CX3CR1",
                       "IFITM1", "ICAM2", "MTSS1", "CDKN1C", "CDH23", "SLC44A2"],  # LILRA/B1 not found in one dataset
        "Monolike": ["FCN1", "HLA-DQA1", "HLA-DQB1", "HSPA1A", "HSPA1B", "CXCR4", "FCGR2A", "FCGR3A"],
        "Ductal": ["KRT19", "SPP1", "CFTR", "CRP", "DEFB1", "CEACAM6", "TSPAN8"],
        "Acinar": ["ALB", "CPA2", "CTRB2", "PNLIP", "CPA1", "PRSS1"],
        "Erythroblast": ["MKI67", "HBA1", "HBB"],
        "NK": ["GNLY", "NKG7", "CD247", "GRIK4", "FCER1G", "TYROBP", "KLRG1", "FCGR3A"],
        "ILC": ["ID2", "PLCG2", "GNLY", "SYNE1"],
        "Endothelial": ["SELE", "FLT1", "KDR", "CD93", "ESAM", "SOX18", "PECAM1", "ESM1", "PLVAP"],
        "Epithelial": ["ACPP", "KLK3", "KLK2", "MSMB", "NEFH", "AZGP1", "RDH11", "PLA2G2A", "TMPRSS2", "NKX3-1"],
        "Fibroblast": ["DCN", "FBLN1", "COL1A2", "IGF1", "C7", "IGFBP5", "CCDC80", "CFD", "LTBP4", "SFRP1"],
        "Naive CD20+ B": ["MS4A1", "IL4R", "IGHD", "FCRL1", "IGHM"],
        "B1 B": ["MS4A1", "SSPN", "ITGB1", "EPHA4", "COL4A4", "PRDM1", "IRF4", "CD38", "XBP1", "PAX5", "BCL11A", "BLK",
                 "IGHD", "IGHM", "ZNF215"],  # Note IGHD and IGHM are negative markers
        "Transitional B": ["MME", "CD38", "CD24", "ACSM3", "MSI2"],
        "Plasma cells": ["MZB1", "HSP90B1", "FNDC3B", "PRDM1", "IGKC", "JCHAIN"],
        "CD4+ T": ["CD4", "IL7R", "TRBC2", "ITGB1", "CCR7", "PTPRC"],
        "CD8+ T": ["CD8A", "CD8B", "GZMK", "GZMA", "CCL5", "GZMB", "GZMH", "GZMA", "PTPRC"],
        "Schwann cell": ["SOX10", "GAP43", "MBP", "MPZ", "NCAM1", "S100A8"]
        # S100 has multiple isoforms, don't know what to put
    }
    return marker_gene_dict


def get_marker_gene_list(dump: bool = False) -> List[str]:
    """If dump is True, dump the list to marker_gene_list.pkl"""
    marker_gene_list = list(itertools.chain.from_iterable(get_marker_gene_dict().values()))
    if dump:
        with open('marker_gene_list.pkl', 'wb') as f:
            pickle.dump(marker_gene_list, f)
    return marker_gene_list


def get_intersection(gene_list: List[str], verbose: bool = True) -> List[str]:
    """Get the intersection of the input gene list and the marker gene list,
    print the markers that did not appear in the list if verbose is True"""
    marker_gene_list = get_marker_gene_list()
    intersection = set(gene_list) & set(marker_gene_list)
    if verbose:
        print('Following markers did not appear in the list:')
        print(set(marker_gene_list) - set(intersection))
    return list(intersection)


def save_marker_expression(adata: AnnData,
                           marker_gene_dict: Dict[str, List[str]] = None,
                           save: bool = True) -> AnnData:
    """Save the expression of marker genes to an h5ad file"""
    intersection = get_intersection(adata.var_names)
    if marker_gene_dict is None:
        marker_gene_dict = get_marker_gene_dict()
    new_adata = adata[:, intersection].copy()
    if save:
        new_adata.write_h5ad('marker_expression.h5ad')
    return new_adata


def check_data_form(mtx):
    """check if the sparse matrix is log transformed or count data (look by eyes)"""
    print('checking')
    import numpy as np
    from scipy.sparse import find
    start_row = 20
    end_row = 30
    start_col = 0
    end_col = 10
    row_indices, col_indices, values = find(mtx[start_row:end_row, start_col:end_col])
    subset_matrix = np.zeros((end_row - start_row, end_col - start_col))
    subset_matrix[row_indices, col_indices] = values
    print(subset_matrix)


def save_fig_to_path(filename, path, file=None):
    path = Path(path)
    if not os.path.exists(path):
        os.makedirs(path)
        print('Created directory:', path)

    if file is None:
        plt.savefig(path / filename)
        plt.close('all')
    else:
        file.savefig(path / filename)
    print(path / filename, 'saved.')


def vis(adata,
        celltype: Union[str, List[str]] = None,
        save: bool = True,
        output: str = 'output',
        marker_file: str = None,
        force_update=False) -> None:
    marker_genes_in_data = dict()
    if marker_file:
        marker_genes = read_marker_from_file(filename=marker_file, force_update=force_update)
    else:
        marker_genes = get_marker_gene_dict()

    for ct, markers in marker_genes.items():
        markers_found = []
        for marker in markers:
            if marker in adata.var_names:
                markers_found.append(marker)
        marker_genes_in_data[ct] = markers_found

    if celltype:
        # robust to celltype being a list or a string
        if isinstance(celltype, str):
            celltype = [celltype]
        # robost to wrong celltype names
        for i in range(len(celltype)):
            if celltype[i] not in marker_genes_in_data.keys():
                print(f'Warning: Cell type "{celltype[i]}" not found in data.')
                celltype[i] = None

        # get the subset of marker_genes_in_data
        temp = {ct: marker_genes_in_data[ct] for ct in celltype if ct is not None}  # dict comprehension
        marker_genes_in_data = temp

    for ct in marker_genes_in_data.keys():
        if len(marker_genes_in_data[ct]) == 0:
            print('No marker genes found for', ct)
            continue
        fig = umap(
            adata,
            color=marker_genes_in_data[ct],
            vmin=0,
            vmax="p99",
            sort_order=False,
            frameon=False,
            cmap="Reds",
            show=False,
        )

        if save:
            save_fig_to_path(f"genes_{ct}.png", output)


def vis_score(adata, save: bool = True, output: str = 'output'):
    adata_score = adata.copy()
    marker_dict = get_marker_gene_dict()
    for ct in marker_dict.keys():
        sc.tl.score_genes(adata_score, marker_dict[ct], score_name='score_' + ct)

    sc.pl.umap(adata_score, color=['score_' + ct for ct in marker_dict.keys()], cmap='Reds', ncols=3, show=False)
    if save:
        save_fig_to_path('scores_for_clusters.png', output)


def annotate(annotation_dict, cell_type, clusters) -> dict:
    """Add annotation to the annotation_dict"""
    if type(clusters) in [int, str]:
        clusters = [str(clusters)]
    for cluster in clusters:
        cluster = str(cluster)
        if cluster in annotation_dict.keys() and annotation_dict[cluster] != cell_type:
            print(f'Warning: {cluster} already annotated as {annotation_dict[cluster]}, now as {cell_type}.')
            continue
        annotation_dict[str(cluster)] = cell_type
    return annotation_dict


def plot_major_lineage(adata, after=False):
    status = 'bbknn_' if after else ''
    markers_dict = {'T': ['PTPRC', 'CD3D', 'CD3E'],
               'M': ['PTPRC', 'MS4A1', 'CST3', 'LYZ', 'ITGAM', 'MRC1'],
               'F': ['COL1A1', 'COL3A1', 'COL1A2', 'RGS5', 'KCNJ8', 'MYH11', 'ACTA2', 'KRT19']}
    for cell_type, markers in markers_dict.items():
        sc.pl.umap(adata, color=get_existing_subset(adata, markers), show=False, save=f'_{status}{cell_type}.pdf', use_raw=True, cmap='Reds')
    gene_list = ["COL1A1", "COL3A1", "COL1A2", "RGS5", "KCNJ8", "MYH11", "ACTA2", "CNN1", "MSLN", "SLPI", "WT1",
                 "CALB2", "KRT19", "EPCAM", "PTPRC", "NKG7", "CD3D", "CD3E", "CD79A", "CD79B", "CST3", "LYZ"]
    key = 'leiden_r2' if after else 'leiden'

    sc.pl.stacked_violin(adata, var_names=get_existing_subset(adata, gene_list), groupby=key, dendrogram=True, use_raw=True,
                         save=f'{status}.pdf')


def get_existing_subset(adata: sc.AnnData, gene_list):
    if adata.raw is None:
        adata.raw = adata
    existing_genes = [gene for gene in gene_list if gene in adata.raw.var_names]
    if len(existing_genes) != len(gene_list):
        print(set(gene_list) - set(existing_genes), 'not found in var_names')
    return existing_genes


def read_gmt(filename, first=None, cache=True):
    """The `first` parameter means how many columns to read, if None, read all columns"""
    with open(filename, 'r') as f:
        lines = f.readlines()
    signatures = {}
    for line in lines:
        line = line.strip()
        line = line.split('\t')
        if first:
            cut = min(first + 2, len(line))
            line = line[: cut]
        signatures[line[0]] = line[2:]
    if cache:
        prefix = filename[:-4]
        pickle.dump(signatures, open(prefix + '.pkl', 'wb'))
    return signatures


def plot_signature_from_dict(adata, signatures):
    for sig in signatures.keys():
        existing_genes = list(set(adata.var_names) & set(signatures[sig]))
        diff = set(existing_genes) - set(signatures[sig])  # why comparing length directly can cause problem?
        if len(diff) != 0:
            print(diff, "not found.")
        sc.pl.umap(adata, color=existing_genes, save=f'_fibro_{sig}.pdf', cmap='Reds', ncols=3)


def prepare_for_umap(adata: sc.AnnData):
    sc.pp.neighbors(adata)
    sc.tl.pca(adata)
    sc.tl.umap(adata)


def find_difference(list1, list2):
    l1 = set(list1)
    l2 = set(list2)
    print("# of unique items in l1:", len(l1))
    print("# of unique items in l2", len(l2))
    s1not2 = l1 - l2
    s2not1 = l2 - l1
    print("In set 1 but not set 2:\n#:",   len(s1not2), "\n", s1not2)
    print("\nIn set 2 but not set 1:\n#:", len(s2not1), "\n", s2not1)


def combine_columns(df: pd.DataFrame, target_col_name: str, source_cols: list):
    cols = [df[col].astype(str) for col in source_cols]
    df[target_col_name] = cols[0]

    for i in range(1, len(cols)):
        df[target_col_name] += "_" + cols[i]