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
2026-09-19 10:58:10 +08:00 · 2020-05-28 18:36:02 -07:00
parent 030eea1898
commit f7585eef1e
5 changed files with 283 additions and 55 deletions
--- 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
--- 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)

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