From f7585eef1e7f264b37858801f037bac629ef2bfe Mon Sep 17 00:00:00 2001 From: bmccandless Date: Thu, 28 May 2020 18:36:02 -0700 Subject: [PATCH] Support for sparse tiledb arrays for the X matrix (#1496) Support for sparse tiledb arrays for the X matrix 1. cxgtool can now output sparse matrices 2. cxg_adaptor and diffexp_cxg updated to handle sparse matrices 3. added a test in test_diffexp to test sparse diffexp and get_X_array --- server/compute/diffexp_cxg.py | 107 ++++++++++++++++++++++++++---- server/converters/cxgtool.py | 115 ++++++++++++++++++++++++++------- server/data_cxg/cxg_adaptor.py | 57 ++++++++++++++-- server/requirements.txt | 1 + server/test/test_diffexp.py | 58 +++++++++++++---- 5 files changed, 283 insertions(+), 55 deletions(-) diff --git a/server/compute/diffexp_cxg.py b/server/compute/diffexp_cxg.py index b1268987..087d9c61 100644 --- a/server/compute/diffexp_cxg.py +++ b/server/compute/diffexp_cxg.py @@ -3,6 +3,8 @@ 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 +from numba import jit + """ See the comments in diffexp_generic for a description of this algorithm @@ -33,26 +35,33 @@ def get_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) - 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) + is_sparse = matrix.schema.sparse - # because all IO is done per-tile, and we are always dense and col-major, + 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 @@ -62,6 +71,7 @@ def diffexp_ttest(adaptor, maskA, maskB, top_n=8, diffexp_lfc_cutoff=0.01): # 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)] @@ -73,10 +83,15 @@ def diffexp_ttest(adaptor, maskA, maskB, top_n=8, diffexp_lfc_cutoff=0.01): 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) - ) + + 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) @@ -111,3 +126,67 @@ def _mean_var_ab(matrix, row_selector_AB, row_selector_A_in_AB, row_selector_B_i 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 diff --git a/server/converters/cxgtool.py b/server/converters/cxgtool.py index 6e5935f4..3290740c 100644 --- a/server/converters/cxgtool.py +++ b/server/converters/cxgtool.py @@ -114,6 +114,13 @@ def main(): help="URL providing more information about the dataset (hint: must be a fully specified absolute URL).", ) parser.add_argument("--out", "--output", "-o", help="output CXG file name") + parser.add_argument( + "--sparse-threshold", + "-s", + type=float, + default=0.0, # force dense by default + help="The X array will be sparse if the percent of non-zeros falls below this value", + ) args = parser.parse_args() global log_level @@ -135,12 +142,15 @@ def main(): obs_names=args.obs_names, about=args.about, extract_colors=not args.disable_custom_colors, + sparse_threshold=args.sparse_threshold, ) log(1, "done") -def write_cxg(adata, container, title, var_names=None, obs_names=None, about=None, extract_colors=False): +def write_cxg( + adata, container, title, var_names=None, obs_names=None, about=None, extract_colors=False, sparse_threshold=5.0 +): if not adata.var.index.is_unique: raise ValueError("Variable index is not unique - unable to convert.") if not adata.obs.index.is_unique: @@ -195,7 +205,7 @@ def write_cxg(adata, container, title, var_names=None, obs_names=None, about=Non log(1, "\t...embeddings created") # X matrix - save_X(container, adata, ctx) + save_X(container, adata, ctx, sparse_threshold) log(1, "\t...X created") @@ -398,44 +408,99 @@ def save_embeddings(container, adata, ctx): log(1, f"\t\t...{name} embedding created") -def create_X(X_name, shape): +def create_X(X_name, shape, is_sparse): """ - Dense, always. Future task: explore if sparse encoding is worth the trouble - below a sparsity threshold. - - The X matrix is access in both row and column oriented patterns, depending on the + The X matrix is accessed in both row and column oriented patterns, depending on the particular operation. Because of the data type, default compression works best. - The tile size (50, 100) and global layout (row/col) was choosen empirically, by benchmarking + The tile size, (50, 100) for dense, and (512,2048) for sparse, + and global layout (row/col) was chosen empirically, by benchmarking the current cellxgene backend. """ filters = tiledb.FilterList([tiledb.ZstdFilter()]) attrs = [tiledb.Attr(dtype=np.float32, filters=filters)] - domain = tiledb.Domain( - tiledb.Dim(name="obs", domain=(0, shape[0] - 1), tile=min(shape[0], 50), dtype=np.uint32), - tiledb.Dim(name="var", domain=(0, shape[1] - 1), tile=min(shape[1], 100), dtype=np.uint32), - ) + if is_sparse: + domain = tiledb.Domain( + tiledb.Dim(name="obs", domain=(0, shape[0] - 1), tile=min(shape[0], 512), dtype=np.uint32), + tiledb.Dim(name="var", domain=(0, shape[1] - 1), tile=min(shape[1], 2048), dtype=np.uint32), + ) + else: + domain = tiledb.Domain( + tiledb.Dim(name="obs", domain=(0, shape[0] - 1), tile=min(shape[0], 50), dtype=np.uint32), + tiledb.Dim(name="var", domain=(0, shape[1] - 1), tile=min(shape[1], 100), dtype=np.uint32), + ) schema = tiledb.ArraySchema( - domain=domain, sparse=False, attrs=attrs, cell_order="row-major", tile_order="col-major" + domain=domain, sparse=is_sparse, attrs=attrs, cell_order="row-major", tile_order="col-major" ) - tiledb.DenseArray.create(X_name, schema) + if is_sparse: + tiledb.SparseArray.create(X_name, schema) + else: + tiledb.DenseArray.create(X_name, schema) -def save_X(container, adata, ctx): +def check_sparse_X(adata, sparse_threshold): + shape = adata.X.shape + stride = min(int(np.power(10, np.around(np.log10(1e9 / shape[1])))), 10_000) + nnz = 0 + maxnnz = int(shape[0] * shape[1] * sparse_threshold / 100) + for row in range(0, shape[0], stride): + lim = min(row + stride, shape[0]) + a = adata.X[row:lim, :] + if type(a) is not np.ndarray: + a = a.toarray() + nnz += np.count_nonzero(a) + if nnz > maxnnz: + return (nnz, lim * shape[1]) + log(2, "\t...rows", row, "to", lim, "nnz", nnz, "nnz percent %5.2f%%" % (100 * nnz / (lim * shape[1]))) + + return (nnz, shape[0] * shape[1]) + + +def save_X(container, adata, ctx, sparse_threshold): # Save X count matrix X_name = f"{container}/X" shape = adata.X.shape - create_X(X_name, shape) + + if sparse_threshold == 100: + is_sparse = True + elif sparse_threshold == 0: + is_sparse = False + else: + nnz, nelem = check_sparse_X(adata, sparse_threshold) + percent = 100.0 * nnz / nelem + if nelem != shape[0] * shape[1]: + log(1, "\t...shape:", str(shape), "non-zeros percent (estimate): %6.2f" % percent) + else: + log(1, "\t...shape:", str(shape), "non-zeros:", nnz, "percent: %6.2f" % percent) + + is_sparse = percent < sparse_threshold + + create_X(X_name, shape, is_sparse) stride = min(int(np.power(10, np.around(np.log10(1e9 / shape[1])))), 10_000) - with tiledb.DenseArray(X_name, mode="w", ctx=ctx) as X: - for row in range(0, shape[0], stride): - lim = min(row + stride, shape[0]) - a = adata.X[row:lim, :] - if type(a) is not np.ndarray: - a = a.toarray() - X[row:lim, :] = a - log(2, "\t...rows", row, "to", lim) - tiledb.consolidate(X_name, ctx=ctx) + if is_sparse: + log(1, "\t...output X as sparse matrix") + with tiledb.SparseArray(X_name, mode="w", ctx=ctx) as X: + for row in range(0, shape[0], stride): + lim = min(row + stride, shape[0]) + a = adata.X[row:lim, :] + if type(a) is not np.ndarray: + a = a.toarray() + indices = np.nonzero(a) + trow = indices[0] + row + X[trow, indices[1]] = a[indices[0], indices[1]] + log(2, "\t...rows", row, "to", lim) + tiledb.consolidate(X_name, ctx=ctx) + + else: + log(1, "\t...output X as dense matrix") + with tiledb.DenseArray(X_name, mode="w", ctx=ctx) as X: + for row in range(0, shape[0], stride): + lim = min(row + stride, shape[0]) + a = adata.X[row:lim, :] + if type(a) is not np.ndarray: + a = a.toarray() + X[row:lim, :] = a + log(2, "\t...rows", row, "to", lim) tiledb.consolidate(X_name, ctx=ctx) diff --git a/server/data_cxg/cxg_adaptor.py b/server/data_cxg/cxg_adaptor.py index accab96a..bbdd94c3 100644 --- a/server/data_cxg/cxg_adaptor.py +++ b/server/data_cxg/cxg_adaptor.py @@ -171,7 +171,11 @@ class CxgAdaptor(DataAdaptor): @staticmethod def _open_array(uri, tiledb_ctx): - return tiledb.DenseArray(uri, mode="r", ctx=tiledb_ctx) + with tiledb.Array(uri, mode="r", ctx=tiledb_ctx) as array: + if array.schema.sparse: + return tiledb.SparseArray(uri, mode="r", ctx=tiledb_ctx) + else: + return tiledb.DenseArray(uri, mode="r", ctx=tiledb_ctx) def open_array(self, name): try: @@ -200,15 +204,58 @@ class CxgAdaptor(DataAdaptor): meta = self.open_array("cxg_group_metadata").meta return json.loads(meta["cxg_category_colors"]) if "cxg_category_colors" in meta else dict() + def __remap_indices(self, coord_range, coord_mask, coord_data): + """ + This function maps the indices in coord_data, which could be in the range [0,coord_range), to + a range that only includes the number of indices encoded in coord_mask. + coord_range is the maxinum size of the range (e.g. get_shape()[0] or get_shape()[1]) + coord_mask is a mask passed into the get_X_array, of size coord_range + coord_data are indices representing locations of non-zero values, in the range [0,coord_range). + + For example, say + coord_mask = [1,0,1,0,0,1] + coord_data = [2,0,2,2,5] + + The function computes the following: + indices = [0,2,5] + ncoord = 3 + maprange = [0,1,2] + mapindex = [0,0,1,0,0,2] + coordindices = [1,0,1,1,2] + """ + if coord_mask is None: + return coord_range, coord_data + + indices = np.where(coord_mask)[0] + ncoord = indices.shape[0] + maprange = np.arange(ncoord) + mapindex = np.zeros(indices[-1] + 1, dtype=int) + mapindex[indices] = maprange + coordindices = mapindex[coord_data] + return ncoord, coordindices + def get_X_array(self, obs_mask=None, var_mask=None): 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[:, :] + + if X.schema.sparse: + if obs_items == slice(None) and var_items == slice(None): + data = X[:, :] + else: + data = X.multi_index[obs_items, var_items] + nrows, obsindices = self.__remap_indices(X.shape[0], obs_mask, data["obs"]) + ncols, varindices = self.__remap_indices(X.shape[1], var_mask, data["var"]) + densedata = np.zeros((nrows, ncols), dtype=self.get_X_array_dtype()) + densedata[obsindices, varindices] = data[""] + return densedata + else: - data = X.multi_index[obs_items, var_items][""] - return data + if obs_items == slice(None) and var_items == slice(None): + data = X[:, :] + else: + data = X.multi_index[obs_items, var_items][""] + return data def get_shape(self): X = self.open_array("X") diff --git a/server/requirements.txt b/server/requirements.txt index f5658bd8..488713a1 100644 --- a/server/requirements.txt +++ b/server/requirements.txt @@ -11,6 +11,7 @@ flask-talisman>=0.7.0 flatbuffers>=1.10.0 flatten-dict>=0.2.0 fsspec>=0.4.4 +numba>=0.49.1 numpy>=1.16.0 packaging>=20.0 pandas>=0.24.2 diff --git a/server/test/test_diffexp.py b/server/test/test_diffexp.py index c75731cb..a5d222f7 100644 --- a/server/test/test_diffexp.py +++ b/server/test/test_diffexp.py @@ -1,11 +1,13 @@ import unittest from server.data_common.matrix_loader import MatrixDataLoader -from server.common.app_config import AppConfig +from server.test import PROJECT_ROOT, app_config import server.compute.diffexp_cxg as diffexp_cxg import server.compute.diffexp_generic as diffexp_generic +from server.converters.cxgtool import write_cxg import numpy as np - -from server.test import PROJECT_ROOT +import tempfile +import anndata +import os class DiffExpTest(unittest.TestCase): @@ -13,12 +15,9 @@ class DiffExpTest(unittest.TestCase): adaptor types and different algorithms.""" def load_dataset(self, path): - app_config = AppConfig() - app_config.single_dataset__datapath = path - app_config.server__verbose = True - app_config.complete_config() + config = app_config(path) loader = MatrixDataLoader(path) - adaptor = loader.open(app_config) + adaptor = loader.open(config) return adaptor def get_mask(self, adaptor, start, stride): @@ -46,9 +45,14 @@ class DiffExpTest(unittest.TestCase): self.assertEqual(len(results), len(expects)) for result, expect in zip(results, expects): self.assertEqual(result[0], expect[0]) - self.assertAlmostEqual(result[1], expect[1]) - self.assertAlmostEqual(result[2], expect[2]) - self.assertAlmostEqual(result[3], expect[3]) + self.assertTrue(np.isclose(result[1], expect[1], 1e-6, 1e-6)) + self.assertTrue(np.isclose(result[2], expect[2], 1e-6, 1e-6)) + self.assertTrue(np.isclose(result[3], expect[3], 1e-6, 1e-6)) + + def get_X_col(self, adaptor, cols): + varmask = np.zeros(adaptor.get_shape()[1], dtype=bool) + varmask[cols] = True + return adaptor.get_X_array(None, varmask) def test_anndata_default(self): """Test an anndata adaptor with its default diffexp algorithm (diffexp_generic)""" @@ -80,3 +84,35 @@ class DiffExpTest(unittest.TestCase): # run it directly results = diffexp_generic.diffexp_ttest(adaptor, maskA, maskB, 10) self.check_1_10_2_10(results) + + def test_cxg_sparse(self): + with tempfile.TemporaryDirectory() as dirname: + sparsename = os.path.join(dirname, "sparse.cxg") + densename = os.path.join(dirname, "dense.cxg") + source_h5ad = anndata.read_h5ad(f"{PROJECT_ROOT}/example-dataset/pbmc3k.h5ad") + # create a cxg sparse array + write_cxg(adata=source_h5ad, container=sparsename, title="pbmc3k", sparse_threshold=100) + write_cxg(adata=source_h5ad, container=densename, title="pbmc3k", sparse_threshold=0) + adaptor_sparse = self.load_dataset(sparsename) + adaptor_dense = self.load_dataset(densename) + + col_results = [] + for adaptor in (adaptor_sparse, adaptor_dense): + maskA = self.get_mask(adaptor, 1, 10) + maskB = self.get_mask(adaptor, 2, 10) + diffexp_results = diffexp_cxg.diffexp_ttest(adaptor, maskA, maskB, 10) + self.check_1_10_2_10(diffexp_results) + topcols = [x[0] for x in diffexp_results] + cols = self.get_X_col(adaptor, topcols) + assert cols.shape[0] == adaptor.get_shape()[0] + assert cols.shape[1] == len(diffexp_results) + col_results.append(cols) + + x = adaptor.get_X_array() + print(x) + + for row in range(col_results[0].shape[0]): + for col in range(col_results[0].shape[1]): + sval = col_results[0][row][col] + dval = col_results[1][row][col] + self.assertTrue(np.isclose(sval, dval, 1e-6, 1e-6))