s2_hvg.py

import argparse
# deal with input
parser = argparse.ArgumentParser(description='Select hvg, perform dimensionality reduction and clustering')
parser.add_argument('-n', '--number', help='the number of hvgs selected', default=2000, type=int)
parser.add_argument('-p', '--prefix', help='prefix of the input file', required=True, type=str)
parser.add_argument('-b', '--batch', help='batch key of the file', default='patient', type=str, required=False)
parser.add_argument('--plot_major_lineage', required=False, default=True, type=bool)

args = parser.parse_args()

import os
import re
import sys
import traceback
import warnings
from pathlib import Path

import matplotlib.pyplot as plt
import numpy as np
import pandas as pd
import scanpy as sc
import seaborn as sns

from util import *

sc.settings.verbosity = 3
warnings.filterwarnings("ignore")
sc.settings.set_figure_params(
    dpi=80,
    facecolor='white',
    frameon=False
)

n = args.number
prefix = args.prefix
batchkey = args.batch

input_file = prefix + '.h5ad'
identifier = ''.join(prefix.split('_')[1:-1])
hvg_file  = Path(f"d3_{identifier}_hvg.h5ad")
umap_file = Path(f"d4_{identifier}_umap.h5ad")

print(f'''
-------------------------
Input file:
{input_file}

Do:
remove doublets
select {n} hvgs

Output file:
{hvg_file}
{umap_file}
-------------------------
      ''')

# read data
try:
    adata = sc.read_h5ad(input_file)
    print(adata)
    if batchkey not in adata.obs.columns:
        print(f'batch key {batchkey} not found in obs, please check')
        sys.exit(1)
    # choose hvg
    adata.uns['log1p']['base'] = None
    adata.obs[batchkey] = adata.obs[batchkey].astype('category')
    sc.pp.highly_variable_genes(adata, n_top_genes=n, flavor='seurat', n_bins=20, inplace=True, batch_key=batchkey) # seurat_v3 is for count data, seurat is for log data
    adata.raw = adata
    adata = adata[:, adata.var['highly_variable']]
    sc.pp.pca(adata, n_comps=50, use_highly_variable=True, svd_solver='arpack')
    adata.write_h5ad(hvg_file)

    print(f'Highly variable genes (n={n}) are saved to {hvg_file}')
    sc.pp.neighbors(adata)
    sc.tl.umap(adata)
    sc.tl.leiden(adata)

    sc.pl.umap(adata, color=['leiden', batchkey], legend_loc='on data', legend_fontsize='small', show=False, save='_batch_before_bbknn.pdf')
    sc.pl.umap(adata, color=['doublet', 'doublet_score'], cmap='Reds', show=False)
    save_fig_to_path('doublet_umap.pdf', 'figures')
    sc.pl.violin(adata, "doublet_score")
    save_fig_to_path('doublet_violin_plot.pdf', 'figures')
    
    adata.write_h5ad(umap_file) # type: ignore

    if args.plot_major_lineage:
        plot_major_lineage(adata)
    print('Finish')
except Exception as e:
    print('Error:\n', e)
    traceback.print_exc(file=sys.stdout)
    print('Current adata object:\n', adata)
    sys.exit(1)