scMultiChat 0.1.0__py3-none-any.whl

MultiChat/Heterogeneous_g_emb/tools/_gene_scores.py

@@ -0,0 +1,346 @@
+"""Predict gene scores based on chromatin accessibility"""
+
+import numpy as np
+import pandas as pd
+import anndata as ad
+import io
+import pybedtools
+from scipy.sparse import (
+    coo_matrix,
+    csr_matrix
+)
+import pkgutil
+
+from ._utils import _uniquify
+
+
+class GeneScores:
+    """A class used to represent gene scores
+
+    Attributes
+    ----------
+
+    Methods
+    -------
+
+    """
+    def __init__(self,
+                 adata,
+                 genome,
+                 gene_anno=None,
+                 tss_upstream=1e5,
+                 tss_downsteam=1e5,
+                 gb_upstream=5000,
+                 cutoff_weight=1,
+                 use_top_pcs=True,
+                 use_precomputed=True,
+                 use_gene_weigt=True,
+                 min_w=1,
+                 max_w=5):
+        """
+        Parameters
+        ----------
+        adata: `Anndata`
+            Input anndata
+        genome : `str`
+            The genome name
+        """
+        self.adata = adata
+        self.genome = genome
+        self.gene_anno = gene_anno
+        self.tss_upstream = tss_upstream
+        self.tss_downsteam = tss_downsteam
+        self.gb_upstream = gb_upstream
+        self.cutoff_weight = cutoff_weight
+        self.use_top_pcs = use_top_pcs
+        self.use_precomputed = use_precomputed
+        self.use_gene_weigt = use_gene_weigt
+        self.min_w = min_w
+        self.max_w = max_w
+
+    def _read_gene_anno(self):
+        """Read in gene annotation
+
+        Parameters
+        ----------
+
+        Returns
+        -------
+
+        """
+        assert (self.genome in ['hg19', 'hg38', 'mm9', 'mm10']),\
+            "`genome` must be one of ['hg19','hg38','mm9','mm10']"
+
+        bin_str = pkgutil.get_data('hge',
+                                   f'data/gene_anno/{self.genome}_genes.bed')
+        gene_anno = pd.read_csv(io.BytesIO(bin_str),
+                                encoding='utf8',
+                                sep='\t',
+                                header=None,
+                                names=['chr', 'start', 'end',
+                                       'symbol', 'strand'])
+        self.gene_anno = gene_anno
+        return self.gene_anno
+
+    def _extend_tss(self, pbt_gene):
+        """Extend transcription start site in both directions
+
+        Parameters
+        ----------
+
+        Returns
+        -------
+
+        """
+        ext_tss = pbt_gene
+        if ext_tss['strand'] == '+':
+            ext_tss.start = max(0, ext_tss.start - self.tss_upstream)
+            ext_tss.end = max(ext_tss.end, ext_tss.start + self.tss_downsteam)
+        else:
+            ext_tss.start = max(0, min(ext_tss.start,
+                                       ext_tss.end - self.tss_downsteam))
+            ext_tss.end = ext_tss.end + self.tss_upstream
+        return ext_tss
+
+    def _extend_genebody(self, pbt_gene):
+        """Extend gene body upstream
+
+        Parameters
+        ----------
+
+        Returns
+        -------
+
+        """
+        ext_gb = pbt_gene
+        if ext_gb['strand'] == '+':
+            ext_gb.start = max(0, ext_gb.start - self.gb_upstream)
+        else:
+            ext_gb.end = ext_gb.end + self.gb_upstream
+        return ext_gb
+
+    def _weight_genes(self):
+        """Weight genes
+
+        Parameters
+        ----------
+
+        Returns
+        -------
+
+        """
+        gene_anno = self.gene_anno
+        gene_size = gene_anno['end'] - gene_anno['start']
+        w = 1/gene_size
+        w_scaled = (self.max_w-self.min_w) * (w-min(w)) / (max(w)-min(w)) \
+            + self.min_w
+        return w_scaled
+
+    def cal_gene_scores(self):
+        """Calculate gene scores
+
+        Parameters
+        ----------
+
+        Returns
+        -------
+
+        """
+        adata = self.adata
+        if self.gene_anno is None:
+            gene_ann = self._read_gene_anno()
+        else:
+            gene_ann = self.gene_anno
+
+        df_gene_ann = gene_ann.copy()
+        df_gene_ann.index = _uniquify(df_gene_ann['symbol'].values)
+        if self.use_top_pcs:
+            mask_p = adata.var['top_pcs']
+        else:
+            mask_p = pd.Series(True, index=adata.var_names)
+        df_peaks = adata.var[mask_p][['chr', 'start', 'end']].copy()
+
+        if 'gene_scores' not in adata.uns_keys():
+            print('Gene scores are being calculated for the first time')
+            print('`use_precomputed` has been ignored')
+            self.use_precomputed = False
+
+        if self.use_precomputed:
+            print('Using precomputed overlap')
+            df_overlap_updated = adata.uns['gene_scores']['overlap'].copy()
+        else:
+            # add the fifth column
+            # so that pybedtool can recognize the sixth column as the strand
+            df_gene_ann_for_pbt = df_gene_ann.copy()
+            df_gene_ann_for_pbt['score'] = 0
+            df_gene_ann_for_pbt = df_gene_ann_for_pbt[['chr', 'start', 'end',
+                                                       'symbol', 'score',
+                                                       'strand']]
+            df_gene_ann_for_pbt['id'] = range(df_gene_ann_for_pbt.shape[0])
+
+            df_peaks_for_pbt = df_peaks.copy()
+            df_peaks_for_pbt['id'] = range(df_peaks_for_pbt.shape[0])
+
+            pbt_gene_ann = pybedtools.BedTool.from_dataframe(
+                df_gene_ann_for_pbt
+                )
+            pbt_gene_ann_ext = pbt_gene_ann.each(self._extend_tss)
+            pbt_gene_gb_ext = pbt_gene_ann.each(self._extend_genebody)
+
+            pbt_peaks = pybedtools.BedTool.from_dataframe(df_peaks_for_pbt)
+
+            # peaks overlapping with extended TSS
+            pbt_overlap = pbt_peaks.intersect(pbt_gene_ann_ext,
+                                              wa=True,
+                                              wb=True)
+            df_overlap = pbt_overlap.to_dataframe(
+                names=[x+'_p' for x in df_peaks_for_pbt.columns]
+                + [x+'_g' for x in df_gene_ann_for_pbt.columns])
+            # peaks overlapping with gene body
+            pbt_overlap2 = pbt_peaks.intersect(pbt_gene_gb_ext,
+                                               wa=True,
+                                               wb=True)
+            df_overlap2 = pbt_overlap2.to_dataframe(
+                names=[x+'_p' for x in df_peaks_for_pbt.columns]
+                + [x+'_g' for x in df_gene_ann_for_pbt.columns])
+
+            # add distance and weight for each overlap
+            df_overlap_updated = df_overlap.copy()
+            df_overlap_updated['dist'] = 0
+
+            for i, x in enumerate(df_overlap['symbol_g'].unique()):
+                # peaks within the extended TSS
+                df_overlap_x = \
+                    df_overlap[df_overlap['symbol_g'] == x].copy()
+                # peaks within the gene body
+                df_overlap2_x = \
+                    df_overlap2[df_overlap2['symbol_g'] == x].copy()
+                # peaks that are not intersecting with the promoter
+                # and gene body of gene x
+                id_overlap = df_overlap_x.index[
+                    ~np.isin(df_overlap_x['id_p'], df_overlap2_x['id_p'])]
+                mask_x = (df_gene_ann['symbol'] == x)
+                range_x = df_gene_ann[mask_x][['start', 'end']].values\
+                    .flatten()
+                if df_overlap_x['strand_g'].iloc[0] == '+':
+                    df_overlap_updated.loc[id_overlap, 'dist'] = pd.concat(
+                        [abs(df_overlap_x.loc[id_overlap, 'start_p']
+                             - (range_x[1])),
+                         abs(df_overlap_x.loc[id_overlap, 'end_p']
+                             - max(0, range_x[0]-self.gb_upstream))],
+                        axis=1, sort=False).min(axis=1)
+                else:
+                    df_overlap_updated.loc[id_overlap, 'dist'] = pd.concat(
+                        [abs(df_overlap_x.loc[id_overlap, 'start_p']
+                             - (range_x[1]+self.gb_upstream)),
+                         abs(df_overlap_x.loc[id_overlap, 'end_p']
+                             - (range_x[0]))],
+                        axis=1, sort=False).min(axis=1)
+
+                n_batch = int(df_gene_ann_for_pbt.shape[0]/5)
+                if i % n_batch == 0:
+                    print(f'Processing: {i/df_gene_ann_for_pbt.shape[0]:.1%}')
+            df_overlap_updated['dist'] = df_overlap_updated['dist']\
+                .astype(float)
+
+            adata.uns['gene_scores'] = dict()
+            adata.uns['gene_scores']['overlap'] = df_overlap_updated.copy()
+
+        df_overlap_updated['weight'] = np.exp(
+            -(df_overlap_updated['dist'].values/self.gb_upstream))
+        mask_w = (df_overlap_updated['weight'] < self.cutoff_weight)
+        df_overlap_updated.loc[mask_w, 'weight'] = 0
+        # construct genes-by-peaks matrix
+        mat_GP = csr_matrix(coo_matrix((df_overlap_updated['weight'],
+                                       (df_overlap_updated['id_g'],
+                                        df_overlap_updated['id_p'])),
+                                       shape=(df_gene_ann.shape[0],
+                                              df_peaks.shape[0])))
+        # adata_GP = ad.AnnData(X=csr_matrix(mat_GP),
+        #                       obs=df_gene_ann,
+        #                       var=df_peaks)
+        # adata_GP.layers['weight'] = adata_GP.X.copy()
+        if self.use_gene_weigt:
+            gene_weights = self._weight_genes()
+            gene_scores = adata[:, mask_p].X * \
+                (mat_GP.T.multiply(gene_weights))
+        else:
+            gene_scores = adata[:, mask_p].X * mat_GP.T
+        adata_CG_atac = ad.AnnData(gene_scores,
+                                   obs=adata.obs.copy(),
+                                   var=df_gene_ann.copy())
+        return adata_CG_atac
+
+
+def gene_scores(adata,
+                genome,
+                gene_anno=None,
+                tss_upstream=1e5,
+                tss_downsteam=1e5,
+                gb_upstream=5000,
+                cutoff_weight=1,
+                use_top_pcs=True,
+                use_precomputed=True,
+                use_gene_weigt=True,
+                min_w=1,
+                max_w=5):
+    """Calculate gene scores
+
+    Parameters
+    ----------
+    adata : AnnData
+        Annotated data matrix.
+    genome : `str`
+        Reference genome. Choose from {'hg19', 'hg38', 'mm9', 'mm10'}
+    gene_anno : `pandas.DataFrame`, optional (default: None)
+        Dataframe of gene annotation.
+        If None, built-in gene annotation will be used depending on `genome`;
+        If provided, custom gene annotation will be used instead.
+    tss_upstream : `int`, optional (default: 1e5)
+        The number of base pairs upstream of TSS
+    tss_downsteam : `int`, optional (default: 1e5)
+        The number of base pairs downstream of TSS
+    gb_upstream : `int`, optional (default: 5000)
+        The number of base pairs upstream by which gene body is extended.
+        Peaks within the extended gene body are given the weight of 1.
+    cutoff_weight : `float`, optional (default: 1)
+        Weight cutoff for peaks
+    use_top_pcs : `bool`, optional (default: True)
+        If True, only peaks associated with top PCs will be used
+    use_precomputed : `bool`, optional (default: True)
+        If True, overlap bewteen peaks and genes
+        (stored in `adata.uns['gene_scores']['overlap']`) will be imported
+    use_gene_weigt : `bool`, optional (default: True)
+        If True, for each gene, the number of peaks assigned to it
+        will be rescaled based on gene size
+    min_w : `int`, optional (default: 1)
+        The minimum weight for each gene.
+        Only valid if `use_gene_weigt` is True
+    max_w : `int`, optional (default: 5)
+        The maximum weight for each gene.
+        Only valid if `use_gene_weigt` is True
+
+    Returns
+    -------
+    adata_new: AnnData
+        Annotated data matrix.
+        Stores #cells x #genes gene score matrix
+
+    updates `adata` with the following fields.
+    overlap: `pandas.DataFrame`, (`adata.uns['gene_scores']['overlap']`)
+        Dataframe of overlap between peaks and genes
+    """
+    GS = GeneScores(adata,
+                    genome,
+                    gene_anno=gene_anno,
+                    tss_upstream=tss_upstream,
+                    tss_downsteam=tss_downsteam,
+                    gb_upstream=gb_upstream,
+                    cutoff_weight=cutoff_weight,
+                    use_top_pcs=use_top_pcs,
+                    use_precomputed=use_precomputed,
+                    use_gene_weigt=use_gene_weigt,
+                    min_w=min_w,
+                    max_w=max_w)
+    adata_CG_atac = GS.cal_gene_scores()
+    return adata_CG_atac

MultiChat/Heterogeneous_g_emb/tools/_general.py

@@ -0,0 +1,71 @@
+"""General-purpose tools"""
+
+import numpy as np
+from sklearn.cluster import KMeans
+
+
+def discretize(adata,
+               layer=None,
+               n_bins=5,
+               max_bins=100):
+    """Discretize continous values
+
+    Parameters
+    ----------
+    adata: AnnData
+        Annotated data matrix.
+    layer: `str`, optional (default: None)
+        The layer used to perform discretization
+    n_bins: `int`, optional (default: 5)
+        The number of bins to produce.
+        It must be smaller than `max_bins`.
+    max_bins: `int`, optional (default: 100)
+        The number of bins used in the initial approximation.
+        i.e. the number of bins to cluster.
+
+    Returns
+    -------
+    updates `adata` with the following fields
+
+    `.layer['hge']` : `array_like`
+        The matrix of discretized values to build HGE(MultiChat) graph.
+    `.uns['disc']` : `dict`
+        `bin_edges`: The edges of each bin.
+        `bin_count`: The number of values in each bin.
+        `hist_edges`: The edges of each bin \
+                      in the initial approximation.
+        `hist_count`: The number of values in each bin \
+                      for the initial approximation.
+    """
+    if layer is None:
+        X = adata.X
+    else:
+        X = adata.layers[layer]
+    nonzero_cont = X.data
+
+    hist_count, hist_edges = np.histogram(
+        nonzero_cont,
+        bins=max_bins,
+        density=False)
+    hist_centroids = (hist_edges[0:-1] + hist_edges[1:])/2
+
+    kmeans = KMeans(n_clusters=n_bins, random_state=2021).fit(
+        hist_centroids.reshape(-1, 1),
+        sample_weight=hist_count)
+    cluster_centers = np.sort(kmeans.cluster_centers_.flatten())
+
+    padding = (hist_edges[-1] - hist_edges[0])/(max_bins*10)
+    bin_edges = np.array(
+        [hist_edges[0]-padding] +
+        list((cluster_centers[0:-1] + cluster_centers[1:])/2) +
+        [hist_edges[-1]+padding])
+    nonzero_disc = np.digitize(nonzero_cont, bin_edges).reshape(-1,)
+    bin_count = np.unique(nonzero_disc, return_counts=True)[1]
+
+    adata.layers['hge'] = X.copy()
+    adata.layers['hge'].data = nonzero_disc
+    adata.uns['disc'] = dict()
+    adata.uns['disc']['bin_edges'] = bin_edges
+    adata.uns['disc']['bin_count'] = bin_count
+    adata.uns['disc']['hist_edges'] = hist_edges
+    adata.uns['disc']['hist_count'] = hist_count

MultiChat/Heterogeneous_g_emb/tools/_integration.py

@@ -0,0 +1,197 @@
+"""Integration across experimental conditions or single cell modalities"""
+
+import numpy as np
+import anndata as ad
+# from sklearn.metrics.pairwise import pairwise_distances
+from sklearn.utils.extmath import randomized_svd
+from scipy.sparse import csr_matrix, find
+
+from ._utils import _knn
+
+
+def infer_edges(adata_ref,
+                adata_query,
+                feature='highly_variable',
+                n_components=20,
+                random_state=42,
+                layer=None,
+                k=20,
+                metric='euclidean',
+                leaf_size=40,
+                **kwargs):
+    """Infer edges between reference and query observations
+
+    Parameters
+    ----------
+    adata_ref: `AnnData`
+        Annotated reference data matrix.
+    adata_query: `AnnData`
+        Annotated query data matrix.
+    feature: `str`, optional (default: None)
+        Feature used for edges inference.
+        The data type of `.var[feature]` needs to be `bool`
+    n_components: `int`, optional (default: 20)
+        The number of components used in `randomized_svd`
+        for comparing reference and query observations
+    random_state: `int`, optional (default: 42)
+        The seed used for truncated randomized SVD
+    n_top_edges: `int`, optional (default: None)
+        The number of edges to keep
+        If specified, `percentile` will be ignored
+    percentile: `float`, optional (default: 0.01)
+        The percentile of edges to keep
+    k: `int`, optional (default: 5)
+        The number of nearest neighbors to consider within each dataset
+    metric: `str`, optional (default: 'euclidean')
+        The metric to use when calculating distance between
+        reference and query observations
+    layer: `str`, optional (default: None)
+        The layer used to perform edge inference
+        If None, `.X` will be used.
+    kwargs:
+        Other keyword arguments are passed down to `randomized_svd()`
+
+    Returns
+    -------
+    adata_ref_query: `AnnData`
+        Annotated relation matrix betwewn reference and query observations
+        Store reference entity as observations and query entity as variables
+    """
+
+    mask_ref = adata_ref.var[feature]
+    feature_ref = adata_ref.var_names[mask_ref]
+    feature_query = adata_query.var_names
+    feature_shared = list(set(feature_ref).intersection(set(feature_query)))
+    print(f'#shared features: {len(feature_shared)}')
+    if layer is None:
+        X_ref = adata_ref[:, feature_shared].X
+        X_query = adata_query[:, feature_shared].X
+    else:
+        X_ref = adata_ref[:, feature_shared].layers[layer]
+        X_query = adata_query[:, feature_shared].layers[layer]
+
+    if any(X_ref.sum(axis=1) == 0) or any(X_query.sum(axis=1) == 0):
+        raise ValueError(
+            f'Some nodes contain zero expressed {feature} features.\n'
+            f'Please try to include more {feature} features.')
+
+    print('Performing randomized SVD ...')
+    mat = X_ref * X_query.T
+    U, Sigma, VT = randomized_svd(mat,
+                                  n_components=n_components,
+                                  random_state=random_state,
+                                  **kwargs)
+    svd_data = np.vstack((U, VT.T))
+    X_svd_ref = svd_data[:U.shape[0], :]
+    X_svd_query = svd_data[-VT.shape[1]:, :]
+    X_svd_ref = X_svd_ref / (X_svd_ref**2).sum(-1, keepdims=True)**0.5
+    X_svd_query = X_svd_query / (X_svd_query**2).sum(-1, keepdims=True)**0.5
+
+    # print('Searching for neighbors within each dataset ...')
+    # knn_conn_ref, knn_dist_ref = _knn(
+    #     X_ref=X_svd_ref,
+    #     k=k,
+    #     leaf_size=leaf_size,
+    #     metric=metric)
+    # knn_conn_query, knn_dist_query = _knn(
+    #     X_ref=X_svd_query,
+    #     k=k,
+    #     leaf_size=leaf_size,
+    #     metric=metric)
+
+    print('Searching for mutual nearest neighbors ...')
+    knn_conn_ref_query, knn_dist_ref_query = _knn(
+        X_ref=X_svd_ref,
+        X_query=X_svd_query,
+        k=k,
+        leaf_size=leaf_size,
+        metric=metric)
+    knn_conn_query_ref, knn_dist_query_ref = _knn(
+        X_ref=X_svd_query,
+        X_query=X_svd_ref,
+        k=k,
+        leaf_size=leaf_size,
+        metric=metric)
+
+    sum_conn_ref_query = knn_conn_ref_query + knn_conn_query_ref.T
+    id_x, id_y, values = find(sum_conn_ref_query > 1)
+    print(f'{len(id_x)} edges are selected')
+    conn_ref_query = csr_matrix(
+        (values*1, (id_x, id_y)),
+        shape=(knn_conn_ref_query.shape))
+    dist_ref_query = csr_matrix(
+        (knn_dist_ref_query[id_x, id_y].A.flatten(), (id_x, id_y)),
+        shape=(knn_conn_ref_query.shape))
+    # it's easier to distinguish zeros (no connection vs zero distance)
+    # using similarity scores
+    sim_ref_query = csr_matrix(
+        (1/(dist_ref_query.data+1), dist_ref_query.nonzero()),
+        shape=(dist_ref_query.shape))  # similarity scores
+
+    # print('Computing similarity scores ...')
+    # dist_ref_query = pairwise_distances(X_svd_ref,
+    #                                     X_svd_query,
+    #                                     metric=metric)
+    # sim_ref_query = 1/(1+dist_ref_query)
+    # # remove low similarity entries to save memory
+    # sim_ref_query = np.where(
+    #     sim_ref_query < np.percentile(sim_ref_query, pct_keep*100),
+    #     0, sim_ref_query)
+    # sim_ref_query = csr_matrix(sim_ref_query)
+
+    adata_ref_query = ad.AnnData(X=sim_ref_query,
+                                 obs=adata_ref.obs,
+                                 var=adata_query.obs)
+    adata_ref_query.layers['hge'] = conn_ref_query
+    adata_ref_query.obsm['svd'] = X_svd_ref
+    # adata_ref_query.obsp['conn'] = knn_conn_ref
+    # adata_ref_query.obsp['dist'] = knn_dist_ref
+    adata_ref_query.varm['svd'] = X_svd_query
+    # adata_ref_query.varp['conn'] = knn_conn_query
+    # adata_ref_query.varp['dist'] = knn_dist_query
+    return adata_ref_query
+
+
+def trim_edges(adata_ref_query,
+               cutoff=None,
+               n_edges=None):
+    """Trim edges based on the similarity scores
+
+    Parameters
+    ----------
+    adata_ref_query: `AnnData`
+         Annotated relation matrix betwewn reference and query observations.
+    n_edges: `int`, optional (default: None)
+        The number of edges to keep
+        If specified, `percentile` will be ignored
+    cutoff: `float`, optional (default: None)
+        The distance cutoff.
+        If None, it will be decided by `n_top_edges`
+        If specified, `n_top_edges` will be ignored
+
+    Returns
+    -------
+    updates `adata_ref_query` with the following field.
+    `.layers['hge']` : `array_like`
+        relation matrix betwewn reference and query observations
+    """
+    sim_ref_query = adata_ref_query.X
+    if cutoff is None:
+        if n_edges is None:
+            raise ValueError('"cutoff" or "n_edges" has to be specified')
+        else:
+            cutoff = \
+                np.partition(sim_ref_query.data,
+                             (sim_ref_query.size-n_edges))[
+                                 sim_ref_query.size-n_edges]
+            # cutoff = \
+            #     np.partition(sim_ref_query.flatten(),
+            #                  (len(sim_ref_query.flatten())-n_edges))[
+            #                      len(sim_ref_query.flatten())-n_edges]
+    id_x, id_y, values = find(sim_ref_query > cutoff)
+
+    print(f'{len(id_x)} edges are selected')
+    conn_ref_query = csr_matrix(
+        (values*1, (id_x, id_y)),
+        shape=(sim_ref_query.shape))
+    adata_ref_query.layers['hge'] = conn_ref_query