mirror of
https://github.com/chanzuckerberg/cellxgene.git
synced 2026-09-16 21:37:59 +08:00
114 lines
4.4 KiB
Python
114 lines
4.4 KiB
Python
import concurrent.futures
|
|
import numpy as np
|
|
from server.compute.diffexp_generic import diffexp_ttest_from_mean_var, mean_var_n
|
|
from server.data_cxg.cxg_util import pack_selector_from_indices
|
|
from server.common.errors import ComputeError
|
|
|
|
"""
|
|
See the comments in diffexp_generic for a description of this algorithm
|
|
|
|
This implementation runs directly in-process. It is multi- threaded, but not particularly scalable.
|
|
Longer term, will likely move to a distributed framework for this.
|
|
|
|
There are currently no global throttles on simultaneous workers.
|
|
"""
|
|
|
|
diffexp_thread_executor = None
|
|
max_workers = None
|
|
target_workunit = None
|
|
|
|
|
|
def set_config(config_max_workers, config_target_workunit):
|
|
global max_workers
|
|
global target_workunit
|
|
max_workers = config_max_workers
|
|
target_workunit = config_target_workunit
|
|
|
|
|
|
def get_thread_executor():
|
|
global diffexp_thread_executor
|
|
if diffexp_thread_executor is None:
|
|
diffexp_thread_executor = concurrent.futures.ThreadPoolExecutor(max_workers=max_workers)
|
|
return diffexp_thread_executor
|
|
|
|
|
|
def diffexp_ttest(adaptor, maskA, maskB, top_n=8, diffexp_lfc_cutoff=0.01):
|
|
row_selector_A = np.where(maskA)[0]
|
|
row_selector_B = np.where(maskB)[0]
|
|
nA = len(row_selector_A)
|
|
nB = len(row_selector_B)
|
|
matrix = adaptor.open_array("X")
|
|
|
|
dtype = matrix.dtype
|
|
cols = matrix.shape[1]
|
|
tile_extent = [dim.tile for dim in matrix.schema.domain]
|
|
|
|
# The rows from both row_selector_A and row_selector_B are gathered at the
|
|
# same time, then the mean and variance are computed by subsetting on that
|
|
# combined submatrix. Combining the gather reduces number of requests/bandwidth
|
|
# to the data source.
|
|
row_selector_AB = np.union1d(row_selector_A, row_selector_B)
|
|
row_selector_A_in_AB = np.in1d(row_selector_AB, row_selector_A, assume_unique=True)
|
|
row_selector_B_in_AB = np.in1d(row_selector_AB, row_selector_B, assume_unique=True)
|
|
row_selector_AB = pack_selector_from_indices(row_selector_AB)
|
|
|
|
# because all IO is done per-tile, and we are always dense and col-major,
|
|
# use the tile column size as the unit of partition. Possibly access
|
|
# more than one column tile at a time based on the target_workunit.
|
|
# Revisit partitioning if we change the X layout, or start using a non-local execution environment
|
|
# which may have other constraints.
|
|
|
|
# TODO: If the number of row selections is large enough, then the cells_per_coltile will exceed
|
|
# the target_workunit. A potential improvement would be to partition by both columns and rows.
|
|
# However partitioning the rows is slightly more complex due to the arbitrary distribution
|
|
# of row selections that are passed into this algorithm.
|
|
cells_per_coltile = (nA + nB) * tile_extent[1]
|
|
cols_per_partition = max(1, int(target_workunit / cells_per_coltile)) * tile_extent[1]
|
|
col_partitions = [(c, min(c + cols_per_partition, cols)) for c in range(0, cols, cols_per_partition)]
|
|
|
|
meanA = np.zeros((cols,), dtype=np.float64)
|
|
varA = np.zeros((cols,), dtype=np.float64)
|
|
meanB = np.zeros((cols,), dtype=np.float64)
|
|
varB = np.zeros((cols,), dtype=np.float64)
|
|
|
|
executor = get_thread_executor()
|
|
futures = []
|
|
for cols in col_partitions:
|
|
futures.append(
|
|
executor.submit(_mean_var_ab, matrix, row_selector_AB, row_selector_A_in_AB, row_selector_B_in_AB, cols)
|
|
)
|
|
|
|
for future in futures:
|
|
# returns tuple: (meanA, varA, meanB, varB, cols)
|
|
try:
|
|
result = future.result()
|
|
part_meanA, part_varA, part_meanB, part_varB, cols = result
|
|
meanA[cols[0] : cols[1]] += part_meanA
|
|
varA[cols[0] : cols[1]] += part_varA
|
|
meanB[cols[0] : cols[1]] += part_meanB
|
|
varB[cols[0] : cols[1]] += part_varB
|
|
except Exception as e:
|
|
for future in futures:
|
|
future.cancel()
|
|
raise ComputeError(str(e))
|
|
|
|
r = diffexp_ttest_from_mean_var(
|
|
meanA.astype(dtype),
|
|
varA.astype(dtype),
|
|
nA,
|
|
meanB.astype(dtype),
|
|
varB.astype(dtype),
|
|
nB,
|
|
top_n,
|
|
diffexp_lfc_cutoff,
|
|
)
|
|
|
|
return r
|
|
|
|
|
|
def _mean_var_ab(matrix, row_selector_AB, row_selector_A_in_AB, row_selector_B_in_AB, col_range):
|
|
X = matrix.multi_index[row_selector_AB, col_range[0] : col_range[1] - 1][""]
|
|
meanA, varA, n = mean_var_n(X[row_selector_A_in_AB])
|
|
meanB, varB, n = mean_var_n(X[row_selector_B_in_AB])
|
|
return (meanA, varA, meanB, varB, col_range)
|