adipose/integrate_sample.py at master · streetslab/adipose · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
import scvi
import scanpy as sc
import argparse
import os
import pandas as pd
import anndata as ad

from  _utils.load_cellbender_output  import anndata_from_h5  # customized function to deal with scalar datasets in order to read cellbender output


def load_h5(h5_file, flag_file=None, extra_filter_file=None):
    """
       extra_filter_file: cellranger's barcode file filtered_feature_bc_matrix/barcodes.tsv.gz
    """
    try:
        adata = sc.read_10x_h5(h5_file, gex_only=True)
    except:
        # In case scanpy can't read cellbender output
        adata = anndata_from_h5(h5_file)

    adata.var_names_make_unique()
    extra_filter_df = None
    if extra_filter_file:
        extra_filter_df = pd.read_csv(extra_filter_file, header=None)
    if flag_file:
        filter_df = pd.read_csv(flag_file, sep=' ', header=0)
        cells = filter_df[(filter_df['dead_cells'] == False) & (filter_df['scds_DropletType'] == 'singlet')]['Barcode']
        adata_cells = adata[adata.obs_names.isin(cells)]
        return adata_cells[adata_cells.obs_names.isin(extra_filter_df[0])]  if extra_filter_file else adata_cells
    return adata[adata.obs_names.isin(extra_filter_df[0])] if extra_filter_file else adata


parser = argparse.ArgumentParser()
parser.add_argument('-f', nargs=1, required=True)
parser.add_argument('-d', nargs=1,  required=True)
parser.add_argument('-o', nargs=1, required=True)
parser.add_argument('-s', nargs='+', required=True)
parser.add_argument('--filter', nargs=1, default=None)
parser.add_argument('--xfilter', nargs=2, default=None)


if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument('-f', nargs=1, required=True)
    parser.add_argument('-d', nargs=1,  required=True)
    parser.add_argument('-o', nargs=1, required=True)
    parser.add_argument('-s', nargs='+', required=True)
    parser.add_argument('--filter', nargs=1, default=None)
    parser.add_argument('--xfilter', nargs=2, default=None)
    args = parser.parse_args()
    #args = parser.parse_args('-d /home/ubuntu/Adipose/Analysis/CellBender_output/  -o SCVI_integration -f output_filtered.h5  -s BOM --filter barcode_qc_flags.csv --xfilter /home/ubuntu/Adipose/Data/arionas/cellranger_outputs filtered_feature_bc_matrix/barcodes.tsv.gz'.split()) #test
    adata_list = {}
    for sample in args.s:
        h5_file=os.path.join(args.d[0], sample, args.f[0])
        flag_file=os.path.join(args.d[0], sample, args.filter[0]) if args.filter else None
        extra_filter_file=os.path.join(args.xfilter[0], sample, args.xfilter[1]) if args.xfilter else None
        adata_list[sample] = load_h5(h5_file, flag_file, extra_filter_file)

    adata_concat = ad.concat(adata_list, join="inner", label="sample")

    sc.pp.filter_genes(adata_concat, min_counts=3)
    adata_concat.raw = adata_concat # keep full dimension safe
    adata_concat.layers['counts'] = adata_concat.X
    sc.pp.highly_variable_genes(
	adata_concat,
	flavor="seurat_v3",
	n_top_genes=2000,
	layer="counts",
	batch_key="sample",
	subset=True)

    scvi.model.SCVI.setup_anndata(adata_concat, layer="counts", batch_key="sample")
    vae = scvi.model.SCVI(adata_concat, n_layers=2, n_latent=30, gene_likelihood="nb")
    vae.train()
    adata_concat.obsm["X_scVI"] = vae.get_latent_representation()

    print("Save latent representation of integrated samples into .h5ad file\n")
    output_file=os.path.join(args.o[0], "adata_scvi_out.h5ad")
    adata_concat.write_h5ad(output_file)