s9_correlation.py

# run: /lustre/home/jydeng/mambaforge/bin/python ../codes/s9_correlation.py -i cell_type -b sample_name -s "Disease
# state"
import argparse

import numpy as np

parser = argparse.ArgumentParser()
parser.add_argument('-f', default='fibro_obs.csv')
parser.add_argument('-m', default='myeloid_obs.csv')
parser.add_argument('-t', default='T_obs.csv')
parser.add_argument('-b', '--sample_name', required=True)
parser.add_argument('-s', '--state', required=True)
parser.add_argument('-i', '--indir', default='data')
parser.add_argument('--levels', nargs='+', required=False)
parser.add_argument('--heatmap',
                    help='the filename of helper file, type "no" if no heatmap is needed, type "all" if no grouping '
                         'is needed',
                    default='sep')
parser.add_argument('--test_type', default='double', choices=['single', 'double'], help='single sided test or two sided test')
parser.add_argument('--outdir', default='celltype_status')
parser.add_argument('--filter_sample', default=0, type=int)
args = parser.parse_args()

import warnings

warnings.filterwarnings('ignore')
import pandas as pd
import matplotlib.pyplot as plt
import seaborn as sns
import sys

from pathlib import Path
from scipy.stats import ranksums
from scipy.stats import mannwhitneyu
from itertools import combinations
from pandas.api.types import CategoricalDtype

global cell_type_col

# how to calculate statistics?
sample_name = args.sample_name
state = args.state

outdir = Path(f'figures_{args.outdir}')
indir = Path(args.indir)
if not indir.exists:
    print(f'"{str(indir.absolute())}" not exists')
    sys.exit(1)

if args.heatmap == 'no':
    plot_heatmap = 0
elif args.heatmap == 'sep':
    plot_heatmap = 1
elif args.heatmap == 'all':
    plot_heatmap = 3
else:
    plot_heatmap = 2
    if not Path(indir / args.heatmap).exists():
        print(str(Path(indir / args.heatmap).absolute()), 'not found')
        sys.exit(1)

if not outdir.exists():
    outdir.mkdir()
    print('Created', str(outdir.absolute()))


def generate_boxplot(data, title, status_order: list, return_df3=False):
    """df3: the wide form, status order should be ordered"""
    if isinstance(data, str):
        df = pd.read_csv(indir / data, index_col=0)
    else:
        df = data
    status_order = CategoricalDtype(categories=status_order, ordered=True)
    cell_type_col = 'meta.cluster.coarse' if 'meta.cluster.coarse' in df.columns else 'majority_voting'
    df2 = df.groupby([args.sample_name, state, cell_type_col])[cell_type_col].count().unstack().fillna(0)
    df_total = df2.sum(axis=1)
    df3 = df2.div(df_total, axis=0)
    df3.columns.name = None
    df3.reset_index(inplace=True)
    df4 = df3.melt(id_vars=[sample_name, state], var_name=cell_type_col, value_name='percentage')
    df5 = df4.merge(df_total.to_frame(name='total'), on=sample_name)
    df5['count'] = (df5['percentage'] * df5['total'])
    df5.fillna(0, inplace=True)
    df5['count'] = df5['count'].astype(int)
    if args.filter_sample != 0:
        print(f'Filtering out samples with less than {args.filter_sample} {title}.')
        before = len(df5[sample_name].unique())
        print(f"sample: {df5['total'] < args.filter_sample} is filtered out." )
        df5 = df5[df5['total'] > args.filter_sample]
        print(f'{len(df5[sample_name].unique())} samples remaining')
    df5[state] = df5[state].astype(status_order)  # set order
    ax = sns.boxplot(x=cell_type_col, y='percentage', hue=state, data=df5, hue_order=levels)
    plt.xticks(rotation=30)
    plt.title(title)
    plt.savefig(outdir / f'boxplot_{title}.pdf')
    plt.close('all')
    ax = sns.boxplot(x=cell_type_col, y='percentage', hue=state, data=df5)
    ax = sns.stripplot(data=df5, x=cell_type_col, y='percentage', hue=state, dodge=True, palette='Set3', legend=False,
                       alpha=0.8, size=3)
    plt.xticks(rotation=30)
    plt.title(title)
    plt.savefig(outdir / f'boxplot_{title}_2.pdf')
    p_values = all_tests(df5, cell_type_col, levels)
    p_values.to_csv(outdir / f'significance_{title}_{args.test_type}.tsv', sep='\t')
    if return_df3:
        df5.to_csv(outdir / f'boxplot_{title}.csv')
        return df3
    return df5


def all_tests(df, cell_type_col, levels):
    all_p_values = []
    for level1, level2 in combinations(levels, 2):
        p_values_df = conduct_test(df, cell_type_col, state, level1, level2)
        all_p_values.append(p_values_df.iloc[:, 1:])
    all_p_values_df = pd.concat(all_p_values, axis=1)
    all_p_values_df.index = p_values_df['Cell Type']
    return all_p_values_df


def conduct_test(df, cell_type_col, state, level1, level2):
    cell_types = df[cell_type_col].unique()
    p_values = {}
    for cell_type in cell_types:
        data1 = df[(df[cell_type_col] == cell_type) & (df[state] == level1)]['percentage']
        data2 = df[(df[cell_type_col] == cell_type) & (df[state] == level2)]['percentage']
        if len(data1) < 2 or len(data2) < 2:
            p_values[cell_type] = np.nan
        else:
            if args.test_type == 'double':
                _, p = ranksums(data1, data2)
            else:
                alternative = 'less' if data1.mean() < data2.mean() else 'greater'
                _, p = mannwhitneyu(data1, data2, alternative=alternative)
            p_values[cell_type] = p
    title = f'{level1}_vs_{level2}'
    p_values_df = pd.DataFrame(list(p_values.items()), columns=['Cell Type', title])
    p_values_df[title + '_significance'] = p_values_df[title] < 0.05
    return p_values_df


def generate_heatmap(df, condition, states):
    if isinstance(states, str):
        states = [states]
    df_sub = df[df[state].isin(states)]
    df_sub.drop(columns=[state], inplace=True)
    df_sub = df_sub.loc[:, df_sub.sum() > 0]
    df_sub.to_csv(outdir / 'heatmap.csv')
    if df_sub.shape[0] < 2:
        print('Too few samples with condition', condition, states)
        return
    sns.clustermap(df_sub.corr(), cmap='vlag', yticklabels=True, xticklabels=True)
    plt.title(f'{condition}: {states}')
    plt.savefig(outdir / f'heatmap_{condition}.pdf')


levels = args.levels

if levels is None:
    print('No `level` information, use fibroblast as reference')
    fibro = pd.read_csv(indir / args.f)
    levels = list(set(fibro[state].dropna()))

for i in range(len(levels)):
    if levels[i].startswith('\\'):
        levels[i] = levels[i][1:]

print('levels:', levels)

if len(levels) != len(set(levels)):
    print("levels must be unique!")
    sys.exit(1)

plt.figure(figsize=(16, 5), dpi=150)
fibro = generate_boxplot(args.f, 'fibro', levels, True)

plt.figure(figsize=(12, 6), dpi=150)
myeloid = generate_boxplot(args.m, 'myeloid', levels, True)

plt.figure(figsize=(10, 6), dpi=150)
T = generate_boxplot(args.t, 'T', levels, True)

myeloid.drop(columns=[state], inplace=True)
T.drop(columns=[state], inplace=True)
# generate heatmap
df = pd.concat([fibro, myeloid, T], axis=1)

df.drop(columns=[sample_name], inplace=True)
df.dropna(inplace=True)

print('df:', df)

if plot_heatmap == 0:
    sys.exit(0)
elif plot_heatmap == 1:
    conditions = dict(zip(levels, levels))
elif plot_heatmap == 3:
    conditions = {'All': levels}
else:
    with open(indir / args.heatmap, 'r') as f:
        conditions = {}
        lines = f.readlines()
        for line in lines:
            elements = line.split(',')
            condition = elements[0]
            conditions[condition] = elements[1].split()

print("heatmap grouping information:", conditions)

for condition, states in conditions.items():
    generate_heatmap(df, condition, states)