mirror of
https://github.com/chanzuckerberg/cellxgene.git
synced 2026-09-15 20:57:56 +08:00
98c2a1359b
* Specialize diffexp for tiledb This patch adds a new diffexp algorithm which is tuned for tiledb. This algorithm was written by Bruce and is adapted here to plug into the current framework. The anndata_adaptor still calls the original algotithm (which was move from diffexp.py to diffexp_generic.py). The cxg_adaptor now calls the new diffexp_tiledb version. Some code is shared between the two. This is part 1 of the diffexp for tiledb. Further tuning and global throttles are still needed. A script to run and time diffexp with various options is also added: test/run_diffexp.py.
118 lines
4.4 KiB
Python
118 lines
4.4 KiB
Python
import os
|
|
import concurrent.futures
|
|
from itertools import repeat
|
|
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
|
|
|
|
"""
|
|
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.
|
|
"""
|
|
|
|
# number of simultaneous workers, per HTTP request
|
|
MAX_WORKERS = 16
|
|
|
|
|
|
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")
|
|
# mean, variance, N - calculate for both selections
|
|
with MyThreadPoolExecutor(max_workers=2) as executor:
|
|
A = executor.submit(mean_var, matrix, row_selector_A)
|
|
B = executor.submit(mean_var, matrix, row_selector_B)
|
|
meanA, varA = A.result()
|
|
meanB, varB = B.result()
|
|
|
|
return diffexp_ttest_from_mean_var(meanA, varA, nA, meanB, varB, nB, top_n, diffexp_lfc_cutoff)
|
|
|
|
|
|
class DispatchMixins:
|
|
def parallel_dispatch(self, fn, *iterables):
|
|
"""
|
|
Dispatch jobs via concurrent.futures.Executor.submit() and wait for their
|
|
completion, return result via future.result(). Primary purpose is to throttle
|
|
dispatch rate so that only 'MAX_JOBS_QUEUE_LENGTH' jobs are running at any
|
|
given time, reducing overall memory footprint.
|
|
"""
|
|
MAX_JOBS_QUEUE_LENGTH = self._max_workers + 4
|
|
|
|
def submit_more_jobs(jobs, active_jobs):
|
|
for job in jobs:
|
|
future = self.submit(fn, *job)
|
|
active_jobs[future] = job
|
|
if len(active_jobs) >= MAX_JOBS_QUEUE_LENGTH:
|
|
return True
|
|
return False
|
|
|
|
def result_iterator(jobs):
|
|
active_jobs = {} # map of future -> args
|
|
try:
|
|
while submit_more_jobs(jobs, active_jobs) or len(active_jobs) > 0:
|
|
for future in concurrent.futures.as_completed(active_jobs.keys()):
|
|
job = active_jobs[future]
|
|
result = future.result()
|
|
# be careful to not retain dangling references
|
|
del active_jobs[future], future
|
|
yield (result, job)
|
|
except Exception as e:
|
|
print(str(e))
|
|
raise
|
|
finally:
|
|
for future in active_jobs.keys():
|
|
future.cancel()
|
|
|
|
return result_iterator(zip(*iterables))
|
|
|
|
|
|
class MyThreadPoolExecutor(concurrent.futures.ThreadPoolExecutor, DispatchMixins):
|
|
pass
|
|
|
|
|
|
def _mean_var(matrix, row_selector, col_range):
|
|
X = matrix.multi_index[row_selector, col_range[0] : col_range[1] - 1][""]
|
|
mean, var, n = mean_var_n(X)
|
|
return (mean, var)
|
|
|
|
|
|
def mean_var(matrix, row_selector):
|
|
"""
|
|
row_selector: list of row indices
|
|
"""
|
|
dtype = matrix.dtype
|
|
rows, cols = matrix.shape
|
|
tile_extent = [dim.tile for dim in matrix.schema.domain]
|
|
|
|
dispatch_func = _mean_var
|
|
row_selector = pack_selector_from_indices(row_selector)
|
|
|
|
# because all IO is done per-tile, and we are always dense and col-major,
|
|
# use the tile column size as the partition. Revisit partitioning if we
|
|
# change the X layout, or start using a non-local execution environment
|
|
# which may have other constraints.
|
|
cols_per_partition = tile_extent[1]
|
|
col_partitions = [(c, min(c + cols_per_partition, cols)) for c in range(0, cols, cols_per_partition)]
|
|
|
|
max_workers = max(1, min(MAX_WORKERS, os.cpu_count())) # throttle max_workers
|
|
|
|
mean = np.zeros((cols,), dtype=np.float64)
|
|
var = np.zeros((cols,), dtype=np.float64)
|
|
dispatch_args = [repeat(matrix), repeat(row_selector), col_partitions]
|
|
with MyThreadPoolExecutor(max_workers=max_workers) as exec:
|
|
for result in exec.parallel_dispatch(dispatch_func, *dispatch_args):
|
|
# returns tuple: (return_val, dispatch_args)
|
|
m, v = result[0]
|
|
cols = result[1][2]
|
|
mean[cols[0] : cols[1]] += m
|
|
var[cols[0] : cols[1]] += v
|
|
del result, m, v
|
|
|
|
return (mean.astype(dtype), var.astype(dtype))
|