mirror of
https://github.com/chanzuckerberg/cellxgene.git
synced 2026-09-16 13:27:59 +08:00
5c0b8c6296
* Improve diffexp for tiledb - The rows from the A and B sets are gathered and processed at the same time. In this way the matrix is only accessed once instead of twice for each tile. - There is now a single thread queue that gets shared between all callers of the diffexp. This will slow down work if diffexp gets too busy. - There is a target_workunit amount of work given to each thread. Previously the workunit was (rows selected * width of tile), which could be small. Now multiple column tiles can be combined into one workunit. If the target is too small then thread and other overheads may reduce performance. If target_workunit is too large then the size of the gathered sub matrix may take up too much memory. - add configuration parameters (max_workers, cpu_multiplier, and target_workunit)
111 lines
4.4 KiB
Python
111 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)
|