import numpy as np
import matplotlib.pyplot as plt
from scipy.sparse import coo_matrix, csr_matrix, diags
from umap.umap_ import fuzzy_simplicial_set
from anndata import AnnData
import scanpy as sc
import scvelo as scv
import pandas as pd
from tqdm.auto import tqdm
import scipy
import os
import sys
from joblib import Parallel, delayed
from tqdm.auto import tqdm
current_path = os.path.dirname(__file__)
src_path = os.path.join(current_path, "..")
sys.path.append(src_path)
from multivelo import mv_logging as logg
from multivelo import settings
current_path = os.path.dirname(__file__)
sys.path.append(current_path)
from pyWNN import pyWNN
def do_count(fastqs, input_loc, output_loc, whitelist_path=None, tech="10XV3", strand=None, threads=8, memory=4):
"""Get spliced and unspliced counts from fastq data.
Makes use of the kallisto-bustools count function:
https://www.kallistobus.tools/kb_usage/kb_count/
kb-python.readthedocs.io/en/latest/autoapi/kb_python/count/index.html
Parameters
----------
fastqs: `List[str]`
The file locations of the fastqs to process.
input_loc: `str`
The folder location of the reference files.
The folder should contain an index file with the name "index.idx", a
transcripts-to-gene file with the name "t2g.txt", a cDNA
transcripts-to-capture file with the name "cdna_t2c.txt", and an
intron transcripts-to-captured file with the name intron_t2c.txt.
output_loc: `str`
The desired folder location of the output of the function.
whitelist_path: `str` (default: `None`)
Path to a barcode whitelist to use to replace the selected technology's
whitelist.
tech: `str` (default: `10XV3`)
The technology used to collect the single-cell data.
strand: `str` (default: `None`)
The strandedness desired to process the data.
threads: `int`
The number of threads to use for parallel processing.
memory: `int`
Maximum memory (in GB) to use while processing.
Returns
-------
adata_count: :class:`~anndata.AnnData`
An AnnData object containing all the spliced and unspliced counts,
as well as associated gene names.
"""
# convert the number of threads and the amount of allocated memory
# into correctly-formatted strings for running kb count
thread_string = str(threads)
memory_string = str(memory) + "G"
# locations of important files
index_loc = input_loc + "/index.idx"
t2g_loc = input_loc + "/t2g.txt"
cdna_t2c = input_loc + "/cdna_t2c.txt"
intron_t2c = input_loc + "/intron_t2c.txt"
# keep the original argv values in case the user specifies it
orig_argv = sys.argv
# assemble the input array to use for kb count
input_array = ["count",
"count",
"-i", index_loc, "-g", t2g_loc,
"-x", tech,
"-o", output_loc,
"-t", thread_string, "-m", memory_string,
"--workflow", "lamanno",
"-c1", cdna_t2c,
"-c2", intron_t2c,
"--h5ad"]
# specify the stranded-ness of the run
if strand is not None:
input_array.append("--strand")
input_array.append(strand)
# specify the whitelist path of the run
if whitelist_path is not None:
input_array.append("-w")
input_array.append(whitelist_path)
# add the fastq's we're doing the run on
for fastq in fastqs:
input_array.append(fastq)
# use our assembled input array as the parameters for kb count
sys.argv = input_array
# run kb count
kbm.main()
# set argv back to its original value
sys.argv = orig_argv
# get the anndata object for
path = output_loc + "/counts_unfiltered"
adata_count = sc.read(path + "/adata.h5ad")
return adata_count
def prepare_gene_mat(var_dict, peaks, gene_mat, adata_atac_X_copy, i):
for peak in peaks:
if peak in var_dict:
peak_index = var_dict[peak]
gene_mat[:, i] += adata_atac_X_copy[:, peak_index]
def aggregate_peaks_10x(adata_atac, peak_annot_file, linkage_file,
peak_dist=10000, min_corr=0.5, gene_body=False,
return_dict=False, parallel=False, n_jobs=1):
"""Peak to gene aggregation.
This function aggregates promoter and enhancer peaks to genes based on the
10X linkage file.
Parameters
----------
adata_atac: :class:`~anndata.AnnData`
ATAC anndata object which stores raw peak counts.
peak_annot_file: `str`
Peak annotation file from 10X CellRanger ARC.
linkage_file: `str`
Peak-gene linkage file from 10X CellRanger ARC. This file stores highly
correlated peak-peak and peak-gene pair information.
peak_dist: `int` (default: 10000)
Maximum distance for peaks to be included for a gene.
min_corr: `float` (default: 0.5)
Minimum correlation for a peak to be considered as enhancer.
gene_body: `bool` (default: `False`)
Whether to add gene body peaks to the associated promoters.
return_dict: `bool` (default: `False`)
Whether to return promoter and enhancer dictionaries.
Returns
-------
A new ATAC anndata object which stores gene aggreagted peak counts.
Additionally, if `return_dict==True`:
A dictionary which stores genes and promoter peaks.
And a dictionary which stores genes and enhancer peaks.
"""
promoter_dict = {}
distal_dict = {}
gene_body_dict = {}
corr_dict = {}
# read annotations
with open(peak_annot_file) as f:
header = next(f)
tmp = header.split('\t')
if len(tmp) == 4:
cellranger_version = 1
elif len(tmp) == 6:
cellranger_version = 2
else:
raise ValueError('Peak annotation file should contain 4 columns '
'(CellRanger ARC 1.0.0) or 6 columns (CellRanger '
'ARC 2.0.0)')
logg.update(f'CellRanger ARC identified as {cellranger_version}.0.0',
v=1)
if cellranger_version == 1:
for line in f:
tmp = line.rstrip().split('\t')
tmp1 = tmp[0].split('_')
peak = f'{tmp1[0]}:{tmp1[1]}-{tmp1[2]}'
if tmp[1] != '':
genes = tmp[1].split(';')
dists = tmp[2].split(';')
types = tmp[3].split(';')
for i, gene in enumerate(genes):
dist = dists[i]
annot = types[i]
if annot == 'promoter':
if gene not in promoter_dict:
promoter_dict[gene] = [peak]
else:
promoter_dict[gene].append(peak)
elif annot == 'distal':
if dist == '0':
if gene not in gene_body_dict:
gene_body_dict[gene] = [peak]
else:
gene_body_dict[gene].append(peak)
else:
if gene not in distal_dict:
distal_dict[gene] = [peak]
else:
distal_dict[gene].append(peak)
else:
for line in f:
tmp = line.rstrip().split('\t')
peak = f'{tmp[0]}:{tmp[1]}-{tmp[2]}'
gene = tmp[3]
dist = tmp[4]
annot = tmp[5]
if annot == 'promoter':
if gene not in promoter_dict:
promoter_dict[gene] = [peak]
else:
promoter_dict[gene].append(peak)
elif annot == 'distal':
if dist == '0':
if gene not in gene_body_dict:
gene_body_dict[gene] = [peak]
else:
gene_body_dict[gene].append(peak)
else:
if gene not in distal_dict:
distal_dict[gene] = [peak]
else:
distal_dict[gene].append(peak)
# read linkages
with open(linkage_file) as f:
for line in f:
tmp = line.rstrip().split('\t')
if tmp[12] == "peak-peak":
peak1 = f'{tmp[0]}:{tmp[1]}-{tmp[2]}'
peak2 = f'{tmp[3]}:{tmp[4]}-{tmp[5]}'
tmp2 = tmp[6].split('><')[0][1:].split(';')
tmp3 = tmp[6].split('><')[1][:-1].split(';')
corr = float(tmp[7])
for t2 in tmp2:
gene1 = t2.split('_')
for t3 in tmp3:
gene2 = t3.split('_')
# one of the peaks is in promoter, peaks belong to the
# same gene or are close in distance
if (((gene1[1] == "promoter") !=
(gene2[1] == "promoter")) and
((gene1[0] == gene2[0]) or
(float(tmp[11]) < peak_dist))):
if gene1[1] == "promoter":
gene = gene1[0]
else:
gene = gene2[0]
if gene in corr_dict:
# peak 1 is in promoter, peak 2 is not in gene
# body -> peak 2 is added to gene 1
if (peak2 not in corr_dict[gene] and
gene1[1] == "promoter" and
(gene2[0] not in gene_body_dict or
peak2 not in gene_body_dict[gene2[0]])):
corr_dict[gene][0].append(peak2)
corr_dict[gene][1].append(corr)
# peak 2 is in promoter, peak 1 is not in gene
# body -> peak 1 is added to gene 2
if (peak1 not in corr_dict[gene] and
gene2[1] == "promoter" and
(gene1[0] not in gene_body_dict or
peak1 not in gene_body_dict[gene1[0]])):
corr_dict[gene][0].append(peak1)
corr_dict[gene][1].append(corr)
else:
# peak 1 is in promoter, peak 2 is not in gene
# body -> peak 2 is added to gene 1
if (gene1[1] == "promoter" and
(gene2[0] not in
gene_body_dict
or peak2 not in
gene_body_dict[gene2[0]])):
corr_dict[gene] = [[peak2], [corr]]
# peak 2 is in promoter, peak 1 is not in gene
# body -> peak 1 is added to gene 2
if (gene2[1] == "promoter" and
(gene1[0] not in
gene_body_dict
or peak1 not in
gene_body_dict[gene1[0]])):
corr_dict[gene] = [[peak1], [corr]]
elif tmp[12] == "peak-gene":
peak1 = f'{tmp[0]}:{tmp[1]}-{tmp[2]}'
tmp2 = tmp[6].split('><')[0][1:].split(';')
gene2 = tmp[6].split('><')[1][:-1]
corr = float(tmp[7])
for t2 in tmp2:
gene1 = t2.split('_')
# peak 1 belongs to gene 2 or are close in distance
# -> peak 1 is added to gene 2
if ((gene1[0] == gene2) or (float(tmp[11]) < peak_dist)):
gene = gene1[0]
if gene in corr_dict:
if (peak1 not in corr_dict[gene] and
gene1[1] != "promoter" and
(gene1[0] not in gene_body_dict or
peak1 not in gene_body_dict[gene1[0]])):
corr_dict[gene][0].append(peak1)
corr_dict[gene][1].append(corr)
else:
if (gene1[1] != "promoter" and
(gene1[0] not in gene_body_dict or
peak1 not in gene_body_dict[gene1[0]])):
corr_dict[gene] = [[peak1], [corr]]
elif tmp[12] == "gene-peak":
peak2 = f'{tmp[3]}:{tmp[4]}-{tmp[5]}'
gene1 = tmp[6].split('><')[0][1:]
tmp3 = tmp[6].split('><')[1][:-1].split(';')
corr = float(tmp[7])
for t3 in tmp3:
gene2 = t3.split('_')
# peak 2 belongs to gene 1 or are close in distance
# -> peak 2 is added to gene 1
if ((gene1 == gene2[0]) or (float(tmp[11]) < peak_dist)):
gene = gene1
if gene in corr_dict:
if (peak2 not in corr_dict[gene] and
gene2[1] != "promoter" and
(gene2[0] not in gene_body_dict or
peak2 not in gene_body_dict[gene2[0]])):
corr_dict[gene][0].append(peak2)
corr_dict[gene][1].append(corr)
else:
if (gene2[1] != "promoter" and
(gene2[0] not in gene_body_dict or
peak2 not in gene_body_dict[gene2[0]])):
corr_dict[gene] = [[peak2], [corr]]
gene_dict = promoter_dict
enhancer_dict = {}
promoter_genes = list(promoter_dict.keys())
logg.update(f'Found {len(promoter_genes)} genes with promoter peaks', 1)
for gene in promoter_genes:
if gene_body: # add gene-body peaks
if gene in gene_body_dict:
for peak in gene_body_dict[gene]:
if peak not in gene_dict[gene]:
gene_dict[gene].append(peak)
enhancer_dict[gene] = []
if gene in corr_dict: # add enhancer peaks
for j, peak in enumerate(corr_dict[gene][0]):
corr = corr_dict[gene][1][j]
if corr > min_corr:
if peak not in gene_dict[gene]:
gene_dict[gene].append(peak)
enhancer_dict[gene].append(peak)
# aggregate to genes
adata_atac_X_copy = adata_atac.X.A
gene_mat = np.zeros((adata_atac.shape[0], len(promoter_genes)))
var_names = adata_atac.var_names.to_numpy()
var_dict = {}
for i, name in enumerate(var_names):
var_dict.update({name: i})
# if we only want to run one job at a time, then no parallelization
# is necessary
if n_jobs == 1:
parallel = False
if parallel:
# if we want to run in parallel, modify the gene_mat variable with
# multiple cores, calling prepare_gene_mat with joblib.Parallel()
Parallel(n_jobs=n_jobs,
require='sharedmem')(
delayed(prepare_gene_mat)(var_dict,
gene_dict[promoter_genes[i]],
gene_mat,
adata_atac_X_copy,
i)for i in tqdm(range(
len(promoter_genes))))
else:
# if we aren't running in parallel, just call prepare_gene_mat
# from a for loop
for i, gene in tqdm(enumerate(promoter_genes),
total=len(promoter_genes)):
prepare_gene_mat(var_dict,
gene_dict[promoter_genes[i]],
gene_mat,
adata_atac_X_copy,
i)
gene_mat[gene_mat < 0] = 0
gene_mat = AnnData(X=csr_matrix(gene_mat))
gene_mat.obs_names = pd.Index(list(adata_atac.obs_names))
gene_mat.var_names = pd.Index(promoter_genes)
gene_mat = gene_mat[:, gene_mat.X.sum(0) > 0]
if return_dict:
return gene_mat, promoter_dict, enhancer_dict
else:
return gene_mat
def tfidf_norm(adata_atac, scale_factor=1e4, copy=False):
"""TF-IDF normalization.
This function normalizes counts in an AnnData object with TF-IDF.
Parameters
----------
adata_atac: :class:`~anndata.AnnData`
ATAC anndata object.
scale_factor: `float` (default: 1e4)
Value to be multiplied after normalization.
copy: `bool` (default: `False`)
Whether to return a copy or modify `.X` directly.
Returns
-------
If `copy==True`, a new ATAC anndata object which stores normalized counts
in `.X`.
"""
npeaks = adata_atac.X.sum(1)
npeaks_inv = csr_matrix(1.0/npeaks)
tf = adata_atac.X.multiply(npeaks_inv)
idf = diags(np.ravel(adata_atac.X.shape[0] / adata_atac.X.sum(0))).log1p()
if copy:
adata_atac_copy = adata_atac.copy()
adata_atac_copy.X = tf.dot(idf) * scale_factor
return adata_atac_copy
else:
adata_atac.X = tf.dot(idf) * scale_factor
def gen_wnn(adata_rna, adata_adt, dims, nn, random_state=0):
"""Computes inputs for KNN smoothing.
This function calculates the nn_idx and nn_dist matrices needed
to run knn_smooth_chrom().
Parameters
----------
adata_rna: :class:`~anndata.AnnData`
RNA anndata object.
adata_atac: :class:`~anndata.AnnData`
ATAC anndata object.
dims: `List[int]`
Dimensions of data for RNA (index=0) and ATAC (index=1)
nn: `int` (default: `None`)
Top N neighbors to extract for each cell in the connectivities matrix.
Returns
-------
nn_idx: `np.darray` (default: `None`)
KNN index matrix of size (cells, k).
nn_dist: `np.darray` (default: `None`)
KNN distance matrix of size (cells, k).
"""
# make a copy of the original adata objects so as to keep them unchanged
rna_copy = adata_rna.copy()
adt_copy = adata_adt.copy()
sc.tl.pca(rna_copy,
n_comps=dims[0],
random_state=np.random.RandomState(seed=42),
use_highly_variable=True) # run PCA on RNA
lsi = scipy.sparse.linalg.svds(adt_copy.X, k=dims[1]) # run SVD on ADT
# get the lsi result
adt_copy.obsm['X_lsi'] = lsi[0]
# add the PCA from adt to rna
rna_copy.obsm['X_rna_pca'] = rna_copy.obsm.pop('X_pca')
rna_copy.obsm['X_adt_lsi'] = adt_copy.obsm['X_lsi']
# run WNN
WNNobj = pyWNN(rna_copy,
reps=['X_rna_pca', 'X_adt_lsi'],
npcs=dims,
n_neighbors=nn,
seed=42)
adata_seurat = WNNobj.compute_wnn(rna_copy)
# get the matrix storing the distances between each cell and its neighbors
cx = scipy.sparse.coo_matrix(adata_seurat.obsp["WNN_distance"])
# the number of cells
cells = adata_seurat.obsp['WNN_distance'].shape[0]
# define the shape of our final results
# and make the arrays that will hold the results
new_shape = (cells, nn)
nn_dist = np.zeros(shape=new_shape)
nn_idx = np.zeros(shape=new_shape)
# new_col defines what column we store data in
# our result arrays
new_col = 0
# loop through the distance matrices
for i, j, v in zip(cx.row, cx.col, cx.data):
# store the distances between neighbor cells
nn_dist[i][new_col % nn] = v
# for each cell's row, store the row numbers of its neighbor cells
# (1-indexing instead of 0- is a holdover from R multimodalneighbors())
nn_idx[i][new_col % nn] = int(j) + 1
new_col += 1
return nn_idx, nn_dist
def knn_smooth_chrom(adata_atac, nn_idx=None, nn_dist=None, conn=None,
n_neighbors=None):
"""KNN smoothing.
This function smooth (impute) the count matrix with k nearest neighbors.
The inputs can be either KNN index and distance matrices or a pre-computed
connectivities matrix (for example in adata_rna object).
Parameters
----------
adata_atac: :class:`~anndata.AnnData`
ATAC anndata object.
nn_idx: `np.darray` (default: `None`)
KNN index matrix of size (cells, k).
nn_dist: `np.darray` (default: `None`)
KNN distance matrix of size (cells, k).
conn: `csr_matrix` (default: `None`)
Pre-computed connectivities matrix.
n_neighbors: `int` (default: `None`)
Top N neighbors to extract for each cell in the connectivities matrix.
Returns
-------
`.layers['Mc']` stores imputed values.
"""
if nn_idx is not None and nn_dist is not None:
if nn_idx.shape[0] != adata_atac.shape[0]:
raise ValueError('Number of rows of KNN indices does not equal to '
'number of observations.')
if nn_dist.shape[0] != adata_atac.shape[0]:
raise ValueError('Number of rows of KNN distances does not equal '
'to number of observations.')
X = coo_matrix(([], ([], [])), shape=(nn_idx.shape[0], 1))
conn, sigma, rho, dists = fuzzy_simplicial_set(X, nn_idx.shape[1],
None, None,
knn_indices=nn_idx-1,
knn_dists=nn_dist,
return_dists=True)
elif conn is not None:
pass
else:
raise ValueError('Please input nearest neighbor indices and distances,'
' or a connectivities matrix of size n x n, with '
'columns being neighbors.'
' For example, RNA connectivities can usually be '
'found in adata.obsp.')
conn = conn.tocsr().copy()
n_counts = (conn > 0).sum(1).A1
if n_neighbors is not None and n_neighbors < n_counts.min():
conn = top_n_sparse(conn, n_neighbors)
conn.setdiag(1)
conn_norm = conn.multiply(1.0 / conn.sum(1)).tocsr()
adata_atac.layers['Mc'] = csr_matrix.dot(conn_norm, adata_atac.X)
adata_atac.obsp['connectivities'] = conn
def calculate_qc_metrics(adata, **kwargs):
"""Basic QC metrics.
This function calculate basic QC metrics with
`scanpy.pp.calculate_qc_metrics`.
Additionally, total counts and the ratio of unspliced and spliced matrices,
as well as the cell cycle scores (with `scvelo.tl.score_genes_cell_cycle`)
will be computed.
Parameters
----------
adata: :class:`~anndata.AnnData`
RNA anndata object. Required fields: `unspliced` and `spliced`.
Additional parameters passed to `scanpy.pp.calculate_qc_metrics`.
Returns
-------
Outputs of `scanpy.pp.calculate_qc_metrics` and
`scvelo.tl.score_genes_cell_cycle`. total_unspliced, total_spliced: `.var`
total counts of unspliced and spliced matrices.
unspliced_ratio: `.var`
ratio of unspliced counts vs (unspliced + spliced counts).
cell_cycle_score: `.var`
cell cycle score difference between G2M_score and S_score.
"""
sc.pp.calculate_qc_metrics(adata, **kwargs)
if 'spliced' not in adata.layers:
raise ValueError('Spliced matrix not found in adata.layers')
if 'unspliced' not in adata.layers:
raise ValueError('Unspliced matrix not found in adata.layers')
logg.update(adata.layers['spliced'].shape, v=1)
total_s = np.nansum(adata.layers['spliced'].toarray(), axis=1)
total_u = np.nansum(adata.layers['unspliced'].toarray(), axis=1)
logg.update(total_u.shape, v=1)
adata.obs['total_unspliced'] = total_u
adata.obs['total_spliced'] = total_s
adata.obs['unspliced_ratio'] = total_u / (total_s + total_u)
scv.tl.score_genes_cell_cycle(adata)
adata.obs['cell_cycle_score'] = (adata.obs['G2M_score']
- adata.obs['S_score'])
def ellipse_fit(adata,
genes,
color_by='quantile',
n_cols=8,
title=None,
figsize=None,
axis_on=False,
pointsize=2,
linewidth=2
):
"""Fit ellipses to unspliced and spliced phase portraits.
This function plots the ellipse fits on the unspliced-spliced phase
portraits.
Parameters
----------
adata: :class:`~anndata.AnnData`
RNA anndata object. Required fields: `Mu` and `Ms`.
genes: `str`, list of `str`
List of genes to plot.
color_by: `str` (default: `quantile`)
Color by the four quantiles based on ellipse fit if `quantile`. Other
common values are leiden, louvain, celltype, etc.
If not `quantile`, the color field must be present in `.uns`, which
can be pre-computed with `scanpy.pl.scatter`.
For `quantile`, red, orange, green, and blue represent quantile left,
top, right, and bottom, respectively.
If `quantile_scores`, `multivelo.compute_quantile_scores` function
must have been run.
n_cols: `int` (default: 8)
Number of columns to plot on each row.
figsize: `tuple` (default: `None`)
Total figure size.
title: `tuple` (default: `None`)
Title of the figure. Default is `Ellipse Fit`.
axis_on: `bool` (default: `False`)
Whether to show axis labels.
pointsize: `float` (default: 2)
Point size for scatter plots.
linewidth: `float` (default: 2)
Line width for ellipse.
"""
by_quantile = color_by == 'quantile'
by_quantile_score = color_by == 'quantile_scores'
if not by_quantile and not by_quantile_score:
types = adata.obs[color_by].cat.categories
colors = adata.uns[f'{color_by}_colors']
gn = len(genes)
if gn < n_cols:
n_cols = gn
fig, axs = plt.subplots(-(-gn // n_cols), n_cols, squeeze=False,
figsize=(2 * n_cols, 2.4 * (-(-gn // n_cols)))
if figsize is None else figsize)
count = 0
for gene in genes:
u = np.array(adata[:, gene].layers['Mu'])
s = np.array(adata[:, gene].layers['Ms'])
row = count // n_cols
col = count % n_cols
non_zero = (u > 0) & (s > 0)
if np.sum(non_zero) < 10:
count += 1
fig.delaxes(axs[row, col])
continue
mean_u, mean_s = np.mean(u[non_zero]), np.mean(s[non_zero])
std_u, std_s = np.std(u[non_zero]), np.std(s[non_zero])
u_ = (u - mean_u)/std_u
s_ = (s - mean_s)/std_s
X = np.reshape(s_[non_zero], (-1, 1))
Y = np.reshape(u_[non_zero], (-1, 1))
# Ax^2 + Bxy + Cy^2 + Dx + Ey + 1 = 0
A = np.hstack([X**2, X * Y, Y**2, X, Y])
b = -np.ones_like(X)
x, res, _, _ = np.linalg.lstsq(A, b)
x = x.squeeze()
A, B, C, D, E = x
good_fit = B**2 - 4*A*C < 0
theta = np.arctan(B/(A - C))/2 \
if x[0] > x[2] \
else np.pi/2 + np.arctan(B/(A - C))/2
good_fit = good_fit & (theta < np.pi/2) & (theta > 0)
if not good_fit:
count += 1
fig.delaxes(axs[row, col])
continue
x_coord = np.linspace((-mean_s)/std_s, (np.max(s)-mean_s)/std_s, 500)
y_coord = np.linspace((-mean_u)/std_u, (np.max(u)-mean_u)/std_u, 500)
X_coord, Y_coord = np.meshgrid(x_coord, y_coord)
Z_coord = (A * X_coord**2 + B * X_coord * Y_coord + C * Y_coord**2 +
D * X_coord + E * Y_coord + 1)
M0 = np.array([
A, B/2, D/2,
B/2, C, E/2,
D/2, E/2, 1,
]).reshape(3, 3)
M = np.array([
A, B/2,
B/2, C,
]).reshape(2, 2)
l1, l2 = np.sort(np.linalg.eigvals(M))
xc = (B*E - 2*C*D)/(4*A*C - B**2)
yc = (B*D - 2*A*E)/(4*A*C - B**2)
slope_major = np.tan(theta)
theta2 = np.pi/2 + theta
slope_minor = np.tan(theta2)
a = np.sqrt(-np.linalg.det(M0)/np.linalg.det(M)/l2)
b = np.sqrt(-np.linalg.det(M0)/np.linalg.det(M)/l1)
xtop = xc + a*np.cos(theta)
ytop = yc + a*np.sin(theta)
xbot = xc - a*np.cos(theta)
ybot = yc - a*np.sin(theta)
xtop2 = xc + b*np.cos(theta2)
ytop2 = yc + b*np.sin(theta2)
xbot2 = xc - b*np.cos(theta2)
ybot2 = yc - b*np.sin(theta2)
mse = res[0] / np.sum(non_zero)
major = lambda x, y: (y - yc) - (slope_major * (x - xc))
minor = lambda x, y: (y - yc) - (slope_minor * (x - xc))
quant1 = (major(s_, u_) > 0) & (minor(s_, u_) < 0)
quant2 = (major(s_, u_) > 0) & (minor(s_, u_) > 0)
quant3 = (major(s_, u_) < 0) & (minor(s_, u_) > 0)
quant4 = (major(s_, u_) < 0) & (minor(s_, u_) < 0)
if (np.sum(quant1 | quant4) < 10) or (np.sum(quant2 | quant3) < 10):
count += 1
continue
if by_quantile:
axs[row, col].scatter(s_[quant1], u_[quant1], s=pointsize,
c='tab:red', alpha=0.6)
axs[row, col].scatter(s_[quant2], u_[quant2], s=pointsize,
c='tab:orange', alpha=0.6)
axs[row, col].scatter(s_[quant3], u_[quant3], s=pointsize,
c='tab:green', alpha=0.6)
axs[row, col].scatter(s_[quant4], u_[quant4], s=pointsize,
c='tab:blue', alpha=0.6)
elif by_quantile_score:
if 'quantile_scores' not in adata.layers:
raise ValueError('Please run multivelo.compute_quantile_scores'
' first to compute quantile scores.')
axs[row, col].scatter(s_, u_, s=pointsize,
c=adata[:, gene].layers['quantile_scores'],
cmap='RdBu_r', alpha=0.7)
else:
for i in range(len(types)):
filt = adata.obs[color_by] == types[i]
axs[row, col].scatter(s_[filt], u_[filt], s=pointsize,
c=colors[i], alpha=0.7)
axs[row, col].contour(X_coord, Y_coord, Z_coord, levels=[0],
colors=('r'), linewidths=linewidth, alpha=0.7)
axs[row, col].scatter([xc], [yc], c='black', s=5, zorder=2)
axs[row, col].scatter([0], [0], c='black', s=5, zorder=2)
axs[row, col].plot([xtop, xbot], [ytop, ybot], color='b',
linestyle='dashed', linewidth=linewidth, alpha=0.7)
axs[row, col].plot([xtop2, xbot2], [ytop2, ybot2], color='g',
linestyle='dashed', linewidth=linewidth, alpha=0.7)
axs[row, col].set_title(f'{gene} {mse:.3g}')
axs[row, col].set_xlabel('s')
axs[row, col].set_ylabel('u')
common_range = [(np.min([(-mean_s)/std_s, (-mean_u)/std_u])
- (0.05*np.max(s)/std_s)),
(np.max([(np.max(s)-mean_s)/std_s,
(np.max(u)-mean_u)/std_u])
+ (0.05*np.max(s)/std_s))]
axs[row, col].set_xlim(common_range)
axs[row, col].set_ylim(common_range)
if not axis_on:
axs[row, col].xaxis.set_ticks_position('none')
axs[row, col].yaxis.set_ticks_position('none')
axs[row, col].get_xaxis().set_visible(False)
axs[row, col].get_yaxis().set_visible(False)
axs[row, col].xaxis.set_ticks_position('none')
axs[row, col].yaxis.set_ticks_position('none')
axs[row, col].set_frame_on(False)
count += 1
for i in range(col+1, n_cols):
fig.delaxes(axs[row, i])
if title is not None:
fig.suptitle(title, fontsize=15)
else:
fig.suptitle('Ellipse Fit', fontsize=15)
fig.tight_layout(rect=[0, 0.1, 1, 0.98])
def compute_quantile_scores(adata,
n_pcs=30,
n_neighbors=30
):
"""Fit ellipses to unspliced and spliced phase portraits and compute
quantile scores.
This function fit ellipses to unspliced-spliced phase portraits. The cells
are split into four groups (quantiles) based on the axes of the ellipse.
Then the function assigns each quantile a score: -3 for left, -1 for top, 1
for right, and 3 for bottom. These gene-specific values are smoothed with a
connectivities matrix. This is similar to the RNA velocity gene time
assignment.
In addition, a 2-bit tuple is assigned to each of the four quantiles, (0,0)
for left, (1,0) for top, (1,1) for right, and (0,1) for bottom. This is to
mimic the distance relationship between quantiles.
Parameters
----------
adata: :class:`~anndata.AnnData`
RNA anndata object. Required fields: `Mu` and `Ms`.
n_pcs: `int` (default: 30)
Number of principal components to compute connectivities.
n_neighbors: `int` (default: 30)
Number of nearest neighbors to compute connectivities.
Returns
-------
quantile_scores: `.layers`
gene-specific quantile scores
quantile_scores_1st_bit, quantile_scores_2nd_bit: `.layers`
2-bit assignment for gene quantiles
quantile_score_sum: `.obs`
aggreagted quantile scores
quantile_genes: `.var`
genes with good quantilty ellipse fits
"""
neighbors = Neighbors(adata)
neighbors.compute_neighbors(n_neighbors=n_neighbors, knn=True, n_pcs=n_pcs)
conn = neighbors.connectivities
conn.setdiag(1)
conn_norm = conn.multiply(1.0 / conn.sum(1)).tocsr()
quantile_scores = np.zeros(adata.shape)
quantile_scores_2bit = np.zeros((adata.shape[0], adata.shape[1], 2))
quantile_gene = np.full(adata.n_vars, False)
quality_gene_idx = []
for idx, gene in enumerate(adata.var_names):
u = np.array(adata[:, gene].layers['Mu'])
s = np.array(adata[:, gene].layers['Ms'])
non_zero = (u > 0) & (s > 0)
if np.sum(non_zero) < 10:
continue
mean_u, mean_s = np.mean(u[non_zero]), np.mean(s[non_zero])
std_u, std_s = np.std(u[non_zero]), np.std(s[non_zero])
u_ = (u - mean_u)/std_u
s_ = (s - mean_s)/std_s
X = np.reshape(s_[non_zero], (-1, 1))
Y = np.reshape(u_[non_zero], (-1, 1))
# Ax^2 + Bxy + Cy^2 + Dx + Ey + 1 = 0
A = np.hstack([X**2, X * Y, Y**2, X, Y])
b = -np.ones_like(X)
x, res, _, _ = np.linalg.lstsq(A, b)
x = x.squeeze()
A, B, C, D, E = x
good_fit = B**2 - 4*A*C < 0
theta = np.arctan(B/(A - C))/2 \
if x[0] > x[2] \
else np.pi/2 + np.arctan(B/(A - C))/2
good_fit = good_fit & (theta < np.pi/2) & (theta > 0)
if not good_fit:
continue
x_coord = np.linspace((-mean_s)/std_s, (np.max(s)-mean_s)/std_s, 500)
y_coord = np.linspace((-mean_u)/std_u, (np.max(u)-mean_u)/std_u, 500)
X_coord, Y_coord = np.meshgrid(x_coord, y_coord)
M = np.array([
A, B/2,
B/2, C,
]).reshape(2, 2)
l1, l2 = np.sort(np.linalg.eigvals(M))
xc = (B*E - 2*C*D)/(4*A*C - B**2)
yc = (B*D - 2*A*E)/(4*A*C - B**2)
slope_major = np.tan(theta)
theta2 = np.pi/2 + theta
slope_minor = np.tan(theta2)
major = lambda x, y: (y - yc) - (slope_major * (x - xc))
minor = lambda x, y: (y - yc) - (slope_minor * (x - xc))
quant1 = (major(s_, u_) > 0) & (minor(s_, u_) < 0)
quant2 = (major(s_, u_) > 0) & (minor(s_, u_) > 0)
quant3 = (major(s_, u_) < 0) & (minor(s_, u_) > 0)
quant4 = (major(s_, u_) < 0) & (minor(s_, u_) < 0)
if (np.sum(quant1 | quant4) < 10) or (np.sum(quant2 | quant3) < 10):
continue
quantile_scores[:, idx:idx+1] = ((-3.) * quant1 + (-1.) * quant2 + 1.
* quant3 + 3. * quant4)
quantile_scores_2bit[:, idx:idx+1, 0] = 1. * (quant1 | quant2)
quantile_scores_2bit[:, idx:idx+1, 1] = 1. * (quant2 | quant3)
quality_gene_idx.append(idx)
quantile_scores = csr_matrix.dot(conn_norm, quantile_scores)
quantile_scores_2bit[:, :, 0] = csr_matrix.dot(conn_norm,
quantile_scores_2bit[:,
:, 0])
quantile_scores_2bit[:, :, 1] = csr_matrix.dot(conn_norm,
quantile_scores_2bit[:,
:, 1])
adata.layers['quantile_scores'] = quantile_scores
adata.layers['quantile_scores_1st_bit'] = quantile_scores_2bit[:, :, 0]
adata.layers['quantile_scores_2nd_bit'] = quantile_scores_2bit[:, :, 1]
quantile_gene[quality_gene_idx] = True
if settings.VERBOSITY >= 1:
perc_good = np.sum(quantile_gene) / adata.n_vars * 100
logg.update(f'{np.sum(quantile_gene)}/{adata.n_vars} - {perc_good:.3g}%'
'genes have good ellipse fits', v=1)
adata.obs['quantile_score_sum'] = \
np.sum(adata[:, quantile_gene].layers['quantile_scores'], axis=1)
adata.var['quantile_genes'] = quantile_gene
def cluster_by_quantile(adata,
plot=False,
n_clusters=None,
affinity='euclidean',
linkage='ward'
):
"""Cluster genes based on 2-bit quantile scores.
This function cluster similar genes based on their 2-bit quantile score
assignments from ellipse fit.
Hierarchical cluster is done with `sklean.cluster.AgglomerativeClustering`.
Parameters
----------
adata: :class:`~anndata.AnnData`
RNA anndata object. Required fields: `Mu` and `Ms`.
plot: `bool` (default: `False`)
Plot the hierarchical clusters.
n_clusters: `int` (default: None)
The number of clusters to keep.
affinity: `str` (default: `euclidean`)
Metric used to compute linkage. Passed to
`sklean.cluster.AgglomerativeClustering`.
linkage: `str` (default: `ward`)
Linkage criterion to use. Passed to
`sklean.cluster.AgglomerativeClustering`.
Returns
-------
quantile_cluster: `.var`
cluster assignments of genes based on quantiles
"""
from sklearn.cluster import AgglomerativeClustering
if 'quantile_scores_1st_bit' not in adata.layers.keys():
raise ValueError("Quantile scores not found. Please run "
"compute_quantile_scores function first.")
quantile_gene = adata.var['quantile_genes']
if plot or n_clusters is None:
cluster = AgglomerativeClustering(distance_threshold=0,
n_clusters=None,
affinity=affinity,
linkage=linkage)
cluster = cluster.fit(np.vstack((adata[:, quantile_gene]
.layers['quantile_scores_1st_bit'],
adata[:, quantile_gene]
.layers['quantile_scores_2nd_bit']))
.transpose())
# https://scikit-learn.org/stable/auto_examples/cluster/plot_agglomerative_dendrogram.html
def plot_dendrogram(model, **kwargs):
from scipy.cluster.hierarchy import dendrogram
counts = np.zeros(model.children_.shape[0])
n_samples = len(model.labels_)
for i, merge in enumerate(model.children_):
current_count = 0
for child_idx in merge:
if child_idx < n_samples:
current_count += 1
else:
current_count += counts[child_idx - n_samples]
counts[i] = current_count
linkage_matrix = np.column_stack([model.children_,
model.distances_,
counts]).astype(float)
dendrogram(linkage_matrix, **kwargs)
plot_dendrogram(cluster, truncate_mode='level', p=5, no_labels=True)
if n_clusters is not None:
n_clusters = int(n_clusters)
cluster = AgglomerativeClustering(n_clusters=n_clusters,
affinity=affinity,
linkage=linkage)
cluster = cluster.fit_predict(np.vstack((adata[:, quantile_gene].layers
['quantile_scores_1st_bit'],
adata[:, quantile_gene].layers
['quantile_scores_2nd_bit']))
.transpose())
quantile_cluster = np.full(adata.n_vars, -1)
quantile_cluster[quantile_gene] = cluster
adata.var['quantile_cluster'] = quantile_cluster
Datasets
Models
