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

@@ -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:

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__()

server/common/default_config.py

@@ -65,6 +65,7 @@ embeddings:
 diffexp:
   enable: true
   lfc_cutoff: 0.01
+  top_n: 10
 
 data_locator:
   s3:

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

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))

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)

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.

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

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

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()