From 98c2a1359b9cfbdbdc0b9a566002627be05eb76d Mon Sep 17 00:00:00 2001 From: bmccandless Date: Sun, 12 Apr 2020 09:56:55 -0700 Subject: [PATCH] Specialize diffexp for tiledb (#1384) * 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. --- server/common/app_config.py | 3 + server/common/constants.py | 3 - server/common/default_config.py | 1 + .../{diffexp.py => diffexp_generic.py} | 106 ++++++++-------- server/compute/diffexp_tiledb.py | 117 ++++++++++++++++++ server/data_anndata/anndata_adaptor.py | 8 ++ server/data_common/data_adaptor.py | 11 +- server/data_cxg/cxg_adaptor.py | 40 ++---- server/data_cxg/cxg_util.py | 37 ++++++ server/test/run_diffexp.py | 88 +++++++++++++ 10 files changed, 328 insertions(+), 86 deletions(-) rename server/compute/{diffexp.py => diffexp_generic.py} (85%) create mode 100644 server/compute/diffexp_tiledb.py create mode 100644 server/data_cxg/cxg_util.py create mode 100644 server/test/run_diffexp.py diff --git a/server/common/app_config.py b/server/common/app_config.py index aeac8dd9..10862e44 100644 --- a/server/common/app_config.py +++ b/server/common/app_config.py @@ -84,6 +84,7 @@ class AppConfig(object): self.diffexp__enable = dc["diffexp"]["enable"] self.diffexp__lfc_cutoff = dc["diffexp"]["lfc_cutoff"] + self.diffexp__top_n = dc["diffexp"]["top_n"] self.data_locator__s3__region_name = dc["data_locator"]["s3"]["region_name"] @@ -189,6 +190,7 @@ class AppConfig(object): context = dict(messagefn=messagefn) self.handle_server(context) + self.handle_adaptor(context) self.handle_data_locator(context) self.handle_adaptor(context) # may depend on data_locator self.handle_presentation(context) @@ -430,6 +432,7 @@ class AppConfig(object): def handle_diffexp(self, context): self.__check_attr("diffexp__enable", bool) self.__check_attr("diffexp__lfc_cutoff", float) + self.__check_attr("diffexp__top_n", int) if self.single_dataset__datapath: with self.matrix_data_cache_manager.data_adaptor(self.single_dataset__datapath, self) as data_adaptor: diff --git a/server/common/constants.py b/server/common/constants.py index c84ddb7c..a883e23a 100644 --- a/server/common/constants.py +++ b/server/common/constants.py @@ -1,9 +1,6 @@ from enum import Enum -DEFAULT_TOP_N = 10 - - class AugmentedEnum(Enum): def __hash__(self): return self.value.__hash__() diff --git a/server/common/default_config.py b/server/common/default_config.py index ee15734c..2a5f1c35 100644 --- a/server/common/default_config.py +++ b/server/common/default_config.py @@ -65,6 +65,7 @@ embeddings: diffexp: enable: true lfc_cutoff: 0.01 + top_n: 10 data_locator: s3: diff --git a/server/compute/diffexp.py b/server/compute/diffexp_generic.py similarity index 85% rename from server/compute/diffexp.py rename to server/compute/diffexp_generic.py index 55ae9cbe..71fe9a59 100644 --- a/server/compute/diffexp.py +++ b/server/compute/diffexp_generic.py @@ -2,46 +2,7 @@ import numpy as np from scipy import sparse, stats -# Convenience function which handles sparse data -def _mean_var_n(X): - """ - Two-pass variance calculation. Numerically (more) stable - than naive methods (and same method used by numpy.var()) - https://en.wikipedia.org/wiki/Algorithms_for_calculating_variance#Two-pass - """ - # fp_err_occurred is a flag indicating that a floating point error - # occured somewhere in our compute. Used to trigger non-finite - # number handling. - fp_err_occurred = False - - def fp_err_set(err, flag): - nonlocal fp_err_occurred - fp_err_occurred = True - - with np.errstate(divide="call", invalid="call", call=fp_err_set): - n = X.shape[0] - if sparse.issparse(X): - mean = X.mean(axis=0).A1 - dfm = X - mean - sumsq = np.sum(np.multiply(dfm, dfm), axis=0).A1 - v = sumsq / (n - 1) - else: - mean = X.mean(axis=0) - dfm = X - mean - sumsq = np.sum(np.multiply(dfm, dfm), axis=0) - v = sumsq / (n - 1) - - if fp_err_occurred: - mean[np.isfinite(mean) == False] = 0 # noqa: E712 - v[np.isfinite(v) == False] = 0 # noqa: E712 - else: - mean[np.isnan(mean)] = 0 - v[np.isnan(v)] = 0 - - return mean, v, n - - -def diffexp_ttest(data, maskA, maskB, top_n=8, diffexp_lfc_cutoff=0.01): +def diffexp_ttest(adaptor, maskA, maskB, top_n=8, diffexp_lfc_cutoff=0.01): """ Return differential expression statistics for top N variables. @@ -59,26 +20,32 @@ def diffexp_ttest(data, maskA, maskB, top_n=8, diffexp_lfc_cutoff=0.01): - p-values adjusted with Bonferroni correction. https://en.wikipedia.org/wiki/Bonferroni_correction - :param data: DataAdaptor instance + :param adaptor: DataAdaptor instance :param maskA: observation selection mask for set 1 :param maskB: observation selection mask for set 2 :param top_n: number of variables to return stats for :param diffexp_lfc_cutoff: minimum :return: for top N genes, [ varindex, logfoldchange, pval, pval_adj ] """ - shape = data.get_shape() - n_obs = shape[0] - n_var = shape[1] - if top_n > n_obs: - top_n = n_obs + + dataA = adaptor.get_X_array(maskA, None) + dataB = adaptor.get_X_array(maskB, None) # mean, variance, N - calculate for both selections - meanA, vA, nA = _mean_var_n(data.get_X_array(maskA, None)) - meanB, vB, nB = _mean_var_n(data.get_X_array(maskB, None)) + meanA, vA, nA = mean_var_n(dataA) + meanB, vB, nB = mean_var_n(dataB) + res = diffexp_ttest_from_mean_var(meanA, vA, nA, meanB, vB, nB, top_n, diffexp_lfc_cutoff) + + return res + + +def diffexp_ttest_from_mean_var(meanA, varA, nA, meanB, varB, nB, top_n, diffexp_lfc_cutoff): + n_var = meanA.shape[0] + top_n = min(top_n, n_var) # variance / N - vnA = vA / min(nA, nB) # overestimate variance, would normally be nA - vnB = vB / min(nA, nB) # overestimate variance, would normally be nB + vnA = varA / min(nA, nB) # overestimate variance, would normally be nA + vnB = varB / min(nA, nB) # overestimate variance, would normally be nB sum_vn = vnA + vnB # degrees of freedom for Welch's t-test @@ -126,3 +93,42 @@ def diffexp_ttest(data, maskA, maskB, top_n=8, diffexp_lfc_cutoff=0.01): # varIndex, logfoldchange, pval, pval_adj result = [[sort_order[i], logfoldchanges_top_n[i], pvals_top_n[i], pvals_adj_top_n[i]] for i in range(top_n)] return result + + +# Convenience function which handles sparse data +def mean_var_n(X): + """ + Two-pass variance calculation. Numerically (more) stable + than naive methods (and same method used by numpy.var()) + https://en.wikipedia.org/wiki/Algorithms_for_calculating_variance#Two-pass + """ + # fp_err_occurred is a flag indicating that a floating point error + # occured somewhere in our compute. Used to trigger non-finite + # number handling. + fp_err_occurred = False + + def fp_err_set(err, flag): + nonlocal fp_err_occurred + fp_err_occurred = True + + with np.errstate(divide="call", invalid="call", call=fp_err_set): + n = X.shape[0] + if sparse.issparse(X): + mean = X.mean(axis=0).A1 + dfm = X - mean + sumsq = np.sum(np.multiply(dfm, dfm), axis=0).A1 + v = sumsq / (n - 1) + else: + mean = X.mean(axis=0) + dfm = X - mean + sumsq = np.sum(np.multiply(dfm, dfm), axis=0) + v = sumsq / (n - 1) + + if fp_err_occurred: + mean[np.isfinite(mean) == False] = 0 # noqa: E712 + v[np.isfinite(v) == False] = 0 # noqa: E712 + else: + mean[np.isnan(mean)] = 0 + v[np.isnan(v)] = 0 + + return mean, v, n diff --git a/server/compute/diffexp_tiledb.py b/server/compute/diffexp_tiledb.py new file mode 100644 index 00000000..11517ee1 --- /dev/null +++ b/server/compute/diffexp_tiledb.py @@ -0,0 +1,117 @@ +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)) diff --git a/server/data_anndata/anndata_adaptor.py b/server/data_anndata/anndata_adaptor.py index 3b8c7212..84fe5976 100644 --- a/server/data_anndata/anndata_adaptor.py +++ b/server/data_anndata/anndata_adaptor.py @@ -15,6 +15,7 @@ from server.common.utils import series_to_schema from server.common.constants import Axis, MAX_LAYOUTS from server.common.errors import PrepareError, DatasetAccessError, FilterError from server.compute.scanpy import scanpy_umap +import server.compute.diffexp_generic as diffexp_generic anndata_version = version.parse(str(anndata.__version__)).release @@ -325,6 +326,13 @@ class AnndataAdaptor(DataAdaptor): schema = {"name": name, "type": "float32", "dims": dims} return (schema, fbs) + def compute_diffexp_ttest(self, maskA, maskB, top_n=None, lfc_cutoff=None): + if top_n is None: + top_n = self.config.diffexp__top_n + if lfc_cutoff is None: + lfc_cutoff = self.config.diffexp__lfc_cutoff + return diffexp_generic.diffexp_ttest(self, maskA, maskB, top_n, lfc_cutoff) + def get_X_array(self, obs_mask=None, var_mask=None): if obs_mask is None: obs_mask = slice(None) diff --git a/server/data_common/data_adaptor.py b/server/data_common/data_adaptor.py index 9c680a37..1e92cfb3 100644 --- a/server/data_common/data_adaptor.py +++ b/server/data_common/data_adaptor.py @@ -5,9 +5,8 @@ import pandas as pd from os.path import basename, splitext from server.data_common.fbs.matrix import encode_matrix_fbs -from server.common.constants import Axis, DEFAULT_TOP_N +from server.common.constants import Axis from server.common.errors import FilterError, JSONEncodingValueError, ExceedsLimitError -from server.compute.diffexp import diffexp_ttest from server.common.utils import jsonify_numpy from server.common.app_config import AppFeature, AppConfig @@ -298,20 +297,24 @@ class DataAdaptor(metaclass=ABCMeta): except (KeyError, IndexError): raise FilterError("Error parsing filter") if top_n is None: - top_n = DEFAULT_TOP_N + top_n = self.config.diffexp__top_n if self.config.exceeds_limit( "diffexp_cellcount_max", np.count_nonzero(obs_mask_A) + np.count_nonzero(obs_mask_B) ): raise ExceedsLimitError("Diffexp request exceeds max cell count limit") - result = diffexp_ttest(self, obs_mask_A, obs_mask_B, top_n, self.config.diffexp__lfc_cutoff) + result = self.compute_diffexp_ttest(obs_mask_A, obs_mask_B, top_n, self.config.diffexp__lfc_cutoff) try: return jsonify_numpy(result) except ValueError: raise JSONEncodingValueError("Error encoding differential expression to JSON") + @abstractmethod + def compute_diffexp_ttest(self, maskA, maskB, top_n, lfc_cutoff): + pass + @staticmethod def normalize_embedding(embedding): """Normalize embedding layout to meet client assumptions. diff --git a/server/data_cxg/cxg_adaptor.py b/server/data_cxg/cxg_adaptor.py index b5d7f409..cd7f5bfc 100644 --- a/server/data_cxg/cxg_adaptor.py +++ b/server/data_cxg/cxg_adaptor.py @@ -7,6 +7,8 @@ from server.common.utils import path_join from server.common.constants import Axis from server.data_common.data_adaptor import DataAdaptor from server.data_common.fbs.matrix import encode_matrix_fbs +from server.data_cxg.cxg_util import pack_selector_from_mask +import server.compute.diffexp_tiledb as diffexp_tiledb from server.common.immutable_kvcache import ImmutableKVCache import tiledb import numpy as np @@ -185,9 +187,16 @@ class CxgAdaptor(DataAdaptor): def compute_embedding(self, method, filter): raise NotImplementedError("CXG does not yet support re-embedding") + def compute_diffexp_ttest(self, maskA, maskB, top_n=None, lfc_cutoff=None): + if top_n is None: + top_n = self.config.diffexp__top_n + if lfc_cutoff is None: + lfc_cutoff = self.config.diffexp__lfc_cutoff + return diffexp_tiledb.diffexp_ttest(self, maskA, maskB, top_n, lfc_cutoff) + def get_X_array(self, obs_mask=None, var_mask=None): - obs_items = self._convert_mask(obs_mask) - var_items = self._convert_mask(var_mask) + obs_items = pack_selector_from_mask(obs_mask) + var_items = pack_selector_from_mask(var_mask) X = self.open_array("X") if obs_items == slice(None) and var_items == slice(None): data = X[:, :] @@ -401,30 +410,3 @@ class CxgAdaptor(DataAdaptor): fbs = encode_matrix_fbs(df, col_idx=df.columns) return fbs - - @staticmethod - def _convert_mask(boolarray): - """Convert an index mask to a list of ranges or indices that can be used in a multi_index.""" - if boolarray is None: - return slice(None) - assert type(boolarray) == np.ndarray - assert (boolarray.dtype) == bool - - selector = np.nonzero(boolarray)[0] - - if len(selector) == 0: - return slice(None) - - result = [] - current = slice(selector[0], selector[0]) - for sel in selector[1:]: - if sel == current.stop + 1: - current = slice(current.start, sel) - else: - result.append(current if current.start != current.stop else current.start) - current = slice(sel, sel) - - if len(result) == 0 or result[-1] != current: - result.append(current if current.start != current.stop else current.start) - - return result diff --git a/server/data_cxg/cxg_util.py b/server/data_cxg/cxg_util.py new file mode 100644 index 00000000..26ef2297 --- /dev/null +++ b/server/data_cxg/cxg_util.py @@ -0,0 +1,37 @@ +import numpy as np + + +def pack_selector_from_mask(boolarray): + """ + pack all contiguous selectors into slices. Remember that + tiledb multi_index requires INCLUSIVE indices. + """ + + if boolarray is None: + return slice(None) + + assert type(boolarray) == np.ndarray + assert boolarray.dtype == bool + + selector = np.nonzero(boolarray)[0] + return pack_selector_from_indices(selector) + + +def pack_selector_from_indices(selector): + + if len(selector) == 0: + return slice(None) + + result = [] + current = slice(selector[0], selector[0]) + for sel in selector[1:]: + if sel == current.stop + 1: + current = slice(current.start, sel) + else: + result.append(current if current.start != current.stop else current.start) + current = slice(sel, sel) + + if len(result) == 0 or result[-1] != current: + result.append(current if current.start != current.stop else current.start) + + return result diff --git a/server/test/run_diffexp.py b/server/test/run_diffexp.py new file mode 100644 index 00000000..62200379 --- /dev/null +++ b/server/test/run_diffexp.py @@ -0,0 +1,88 @@ +import sys +import argparse +import random +import time +import numpy as np + +import server.compute.diffexp_tiledb as diffexp_tiledb +import server.compute.diffexp_generic as diffexp_generic + +from server.common.app_config import AppConfig +from server.data_common.matrix_loader import MatrixDataLoader +from server.data_cxg.cxg_adaptor import CxgAdaptor + + +def main(): + parser = argparse.ArgumentParser("A command to test diffexp") + parser.add_argument("dataset", help="name of a dataset to load") + parser.add_argument("-na", "--numA", type=int, required=True, help="number of rows in group A") + parser.add_argument("-nb", "--numB", type=int, required=True, help="number of rows in group B") + parser.add_argument("-t", "--trials", default=1, type=int, help="number of trials") + parser.add_argument( + "-a", "--alg", choices=("default", "generic", "tiledb"), default="default", help="algorithm to use" + ) + parser.add_argument("-s", "--show", default=False, action="store_true", help="show the results") + parser.add_argument( + "-n", "--new-selection", default=False, action="store_true", help="change the selection between each trial" + ) + parser.add_argument("--seed", default=1, type=int, help="set the random seed") + + args = parser.parse_args() + + app_config = AppConfig() + app_config.data_locator__s3__region_name = "us-west-2" + app_config.adaptor__cxg_adaptor__tiledb_ctx["vfs.s3.region"] = "us-west-2" + app_config.single_dataset__datapath = args.dataset + app_config.server__verbose = True + app_config.complete_config() + + loader = MatrixDataLoader(args.dataset) + adaptor = loader.open(app_config) + + if args.show: + if isinstance(adaptor, CxgAdaptor): + adaptor.open_array("X").schema.dump() + + numA = args.numA + numB = args.numB + rows = adaptor.get_shape()[0] + + random.seed(args.seed) + + if not args.new_selection: + samples = random.sample(range(rows), numA + numB) + filterA = samples[:numA] + filterB = samples[numA:] + + for i in range(args.trials): + if args.new_selection: + samples = random.sample(range(rows), numA + numB) + filterA = samples[:numA] + filterB = samples[numA:] + + maskA = np.zeros(rows, dtype=bool) + maskA[filterA] = True + maskB = np.zeros(rows, dtype=bool) + maskB[filterB] = True + + t1 = time.time() + if args.alg == "default": + results = adaptor.compute_diffexp_ttest(maskA, maskB) + elif args.alg == "generic": + results = diffexp_generic.diffexp_ttest(adaptor, maskA, maskB) + elif args.alg == "tiledb": + if not isinstance(adaptor, CxgAdaptor): + print("tiledb only works with CxgAdaptor") + sys.exit(1) + results = diffexp_tiledb.diffexp_ttest(adaptor, maskA, maskB) + + t2 = time.time() + print("TIME=", t2 - t1) + + if args.show: + for res in results: + print(res) + + +if __name__ == "__main__": + main()