Refactor czi_hosted and server into backend directory, pull common code into backend/common, refactor tests (#2102)

* move local_server -> backend/server server-> backend/czi_hosted, pull common code into backend/common update imports, tests and make commands

backend/czi_hosted/compute/diffexp_cxg.py

@@ -0,0 +1,199 @@
+import concurrent.futures
+import numpy as np
+
+from numba import jit
+
+from backend.czi_hosted.data_cxg.cxg_util import pack_selector_from_indices
+from backend.czi_hosted.compute.diffexp_generic import diffexp_ttest_from_mean_var, mean_var_n
+from backend.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):
+
+    matrix = adaptor.open_array("X")
+    row_selector_A = np.where(maskA)[0]
+    row_selector_B = np.where(maskB)[0]
+    nA = len(row_selector_A)
+    nB = len(row_selector_B)
+
+    dtype = matrix.dtype
+    cols = matrix.shape[1]
+    tile_extent = [dim.tile for dim in matrix.schema.domain]
+
+    is_sparse = matrix.schema.sparse
+
+    if is_sparse:
+        row_selector_A = pack_selector_from_indices(row_selector_A)
+        row_selector_B = pack_selector_from_indices(row_selector_B)
+    else:
+        # 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 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 = []
+
+    if is_sparse:
+        for cols in col_partitions:
+            futures.append(executor.submit(_mean_var_sparse_ab, matrix, row_selector_A, nA, row_selector_B, nB, cols))
+    else:
+        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))
+
+    if is_sparse:
+        if adaptor.has_array("X_col_shift"):
+            X_col_shift = adaptor.open_array("X_col_shift")[:]
+            meanA += X_col_shift
+            meanB += X_col_shift
+
+    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)
+
+
+def _mean_var_sparse_ab(matrix, row_selector_A, nrows_A, row_selector_B, nrows_B, col_range):
+    meanA, varA = _mean_var_sparse(matrix, row_selector_A, nrows_A, col_range)
+    meanB, varB = _mean_var_sparse(matrix, row_selector_B, nrows_B, col_range)
+    return (meanA, varA, meanB, varB, col_range)
+
+
+@jit(nopython=True)
+def _mean_var_sparse_numba(x, var, nrows, ncols):
+    """Kernel to compute the mean and variance.  It was not clear if this function
+    could be written using numpy, thus avoiding the loops.  Therefore numba is
+    used here to speed things up.  With numba, this function takes a negligible amount
+    of time compared to reading in the sparse matrix"""
+    mean = np.zeros((ncols,), dtype=np.float64)
+    for col, val in zip(var, x):
+        mean[col] += val
+    mean /= nrows
+
+    # optimize the sumsq computation.
+    # since most entries in a sparse matrix are 0, then start by assuming
+    # all values are 0, so fill the sumsq array with nrows * (0 - mean)**2.
+    # as non-zero values are encountered, subtract off the (mean*mean) value
+    # and replace with (val-mean)**2.  Simplifying the expression
+    # gives the following code.
+    sumsq = nrows * np.multiply(mean, mean)
+    for col, val in zip(var, x):
+        sumsq[col] += val * (val - 2 * mean[col])
+    v = sumsq / (nrows - 1)
+    return mean, v
+
+
+def _mean_var_sparse(matrix, selector, nrows, col_range):
+    data = matrix.multi_index[selector, col_range[0] : col_range[1] - 1]
+    x = data[""]
+
+    # tiledb < 0.6.0 and >= 0.6.0 have slightly different interfaces.
+    # the following takes care of both cases:
+    #   older:  data["coords]["var"]
+    #   newer:  data["var"]
+    var = data.get("coords", data)["var"]
+
+    # shift the column indices to start at 0, this
+    # will become the index into the mean and var arrays.
+    var -= col_range[0]
+
+    fp_err_occurred = False
+
+    def fp_err_set(err, flag):
+        nonlocal fp_err_occurred
+        fp_err_occurred = True
+
+    ncols = col_range[1] - col_range[0]
+    with np.errstate(divide="call", invalid="call", call=fp_err_set):
+        mean, v = _mean_var_sparse_numba(x, var, nrows, ncols)
+
+    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

backend/czi_hosted/compute/diffexp_generic.py

@@ -0,0 +1,134 @@
+import numpy as np
+from scipy import sparse, stats
+
+
+def diffexp_ttest(adaptor, maskA, maskB, top_n=8, diffexp_lfc_cutoff=0.01):
+    """
+    Return differential expression statistics for top N variables.
+
+    Algorithm:
+    - compute log fold change (log2(meanA/meanB))
+    - compute Welch's t-test statistic and pvalue (w/ Bonferroni correction)
+    - return top N abs(logfoldchange) where lfc > diffexp_lfc_cutoff
+
+    If there are not N which meet criteria, augment by removing the logfoldchange
+    threshold requirement.
+
+    Notes on alogrithm:
+    - Welch's ttest provides basic statistics test.
+      https://en.wikipedia.org/wiki/Welch%27s_t-test
+    - p-values adjusted with Bonferroni correction.
+      https://en.wikipedia.org/wiki/Bonferroni_correction
+
+    :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 ]
+    """
+
+    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(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 = 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
+    with np.errstate(divide="ignore", invalid="ignore"):
+        dof = sum_vn ** 2 / (vnA ** 2 / (nA - 1) + vnB ** 2 / (nB - 1))
+    dof[np.isnan(dof)] = 1
+
+    # Welch's t-test score calculation
+    with np.errstate(divide="ignore", invalid="ignore"):
+        tscores = (meanA - meanB) / np.sqrt(sum_vn)
+    tscores[np.isnan(tscores)] = 0
+
+    # p-value
+    pvals = stats.t.sf(np.abs(tscores), dof) * 2
+    pvals_adj = pvals * n_var
+    pvals_adj[pvals_adj > 1] = 1  # cap adjusted p-value at 1
+
+    # logfoldchanges: log2(meanA / meanB)
+    logfoldchanges = np.log2(np.abs((meanA + 1e-9) / (meanB + 1e-9)))
+
+    # find all with lfc > cutoff
+    lfc_above_cutoff_idx = np.nonzero(np.abs(logfoldchanges) > diffexp_lfc_cutoff)[0]
+    stats_to_sort = np.abs(tscores)
+
+    # derive sort order
+    if lfc_above_cutoff_idx.shape[0] > top_n:
+        # partition top N
+        rel_t_partition = np.argpartition(stats_to_sort[lfc_above_cutoff_idx], -top_n)[-top_n:]
+        t_partition = lfc_above_cutoff_idx[rel_t_partition]
+        # sort the top N partition
+        rel_sort_order = np.argsort(stats_to_sort[t_partition])[::-1]
+        sort_order = t_partition[rel_sort_order]
+    else:
+        # partition and sort top N, ignoring lfc cutoff
+        partition = np.argpartition(stats_to_sort, -top_n)[-top_n:]
+        rel_sort_order = np.argsort(stats_to_sort[partition])[::-1]
+        indices = np.indices(stats_to_sort.shape)[0]
+        sort_order = indices[partition][rel_sort_order]
+
+    # top n slice based upon sort order
+    logfoldchanges_top_n = logfoldchanges[sort_order]
+    pvals_top_n = pvals[sort_order]
+    pvals_adj_top_n = pvals_adj[sort_order]
+
+    # 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

backend/czi_hosted/compute/scanpy.py

@@ -0,0 +1,53 @@
+import importlib
+import numpy as np
+
+"""
+Wrapper for various scanpy modules.  Will raise NotImplementedError if the scanpy
+module is not installed/available
+"""
+
+
+def get_scanpy_module():
+    try:
+        sc = importlib.import_module("scanpy")
+        # Future: we could enforce versions here, eg, lookat sc.__version__
+        return sc
+    except ModuleNotFoundError as e:
+        raise NotImplementedError("Please install scanpy to enable UMAP re-embedding") from e
+    except Exception as e:
+        # will capture other ImportError corner cases
+        raise NotImplementedError() from e
+
+
+def scanpy_umap(adata, obs_mask=None, pca_options={}, neighbors_options={}, umap_options={}):
+    """
+    Given adata and an obs mask, return a new embedding for adata[obs_mask, :]
+    as an ndarray of shape (len(obs_mask), N), where N>=2.
+
+    Do NOT mutate adata.
+    """
+
+    # backed mode is incompatible with the current implementation
+    if adata.isbacked:
+        raise NotImplementedError("Backed mode is incompatible with re-embedding")
+
+    # safely get scanpy module, which may not be present.
+    sc = get_scanpy_module()
+
+    # https://github.com/theislab/anndata/issues/311
+    obs_mask = slice(None) if obs_mask is None else obs_mask
+    adata = adata[obs_mask, :].copy()
+
+    for k in list(adata.obsm.keys()):
+        del adata.obsm[k]
+    for k in list(adata.uns.keys()):
+        del adata.uns[k]
+
+    sc.pp.pca(adata, zero_center=None, n_comps=min(adata.n_vars - 1, 50), **pca_options)
+    sc.pp.neighbors(adata, **neighbors_options)
+    sc.tl.umap(adata, **umap_options)
+
+    umap = adata.obsm["X_umap"]
+    result = np.full((obs_mask.shape[0], umap.shape[1]), np.NaN)
+    result[obs_mask] = umap
+    return result