Fountain/plot.py at main · BioX-NKU/Fountain · 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
85
86
87
88
89
#!/usr/bin/env python

import numpy as np
import scanpy as sc
import matplotlib.pyplot as plt
import seaborn as sns
import anndata as  ad


def embedding(
        adata,
        color='celltype',
        color_map=None,
        groupby='batch',
        groups=None,
        cond2=None,
        v2=None,
        save=None,
        legend_loc='right margin',
        legend_fontsize=None,
        legend_fontweight='bold',
        sep='_',
        basis='X_umap',
        size=10,
        show=True,
    ):
    """
    plot separated embeddings with others as background

    Parameters
    ----------
    adata
        AnnData
    color
        meta information to be shown
    color_map
        specific color map
    groupby
        condition which is based-on to separate
    groups
        specific groups to be shown
    cond2
        another targeted condition
    v2
        another targeted values of another condition
    basis
        embeddings used to visualize, default is X_umap for UMAP
    size
        dot size on the embedding
    """

    if groups is None:
        groups = adata.obs[groupby].cat.categories


    for b in groups:
        adata.obs['tmp'] = adata.obs[color].astype(str)
        adata.obs['tmp'][adata.obs[groupby]!=b] = ''
        if cond2 is not None:
            adata.obs['tmp'][adata.obs[cond2]!=v2] = ''
            groups = list(adata[(adata.obs[groupby]==b) &
                                (adata.obs[cond2]==v2)].obs[color].astype('category').cat.categories.values)
            size = min(size, 120000/len(adata[(adata.obs[groupby]==b) & (adata.obs[cond2]==v2)]))
        else:
            groups = list(adata[adata.obs[groupby]==b].obs[color].astype('category').cat.categories.values)
            size = min(size, 120000/len(adata[adata.obs[groupby]==b]))
        adata.obs['tmp'] = adata.obs['tmp'].astype('category')
        if color_map is not None:
            palette = [color_map[i] if i in color_map else 'gray' for i in adata.obs['tmp'].cat.categories]
        else:
            palette = None

        title = b if cond2 is None else v2+sep+b
        if save is not None:
            save_ = '_'+b+save
            show = False
        else:
            save_ = None
            show = True
        sc.pl.embedding(adata, color='tmp', basis=basis, groups=groups, title=title, palette=palette, size=size, save=save_,
                   legend_loc=legend_loc, legend_fontsize=legend_fontsize, legend_fontweight=legend_fontweight, show=show)
        del adata.obs['tmp']
        del adata.uns['tmp_colors']