mirror of
https://github.com/chanzuckerberg/cellxgene.git
synced 2026-10-04 06:38:11 +08:00
sparse column shift encoding. (#1502)
Many of our matrices are log normalized, which tends to eliminate the number of non zero values (if there were any). This prevents the matrix from being stored as a sparse matrix. The solution here is to use a simple transformation to make it sparse again. The most common value from each column is subtracted from that column. These values that were subtracted are saved in an array called X_col_shift. The cellxgene code needs to understand how to undo the transformation when operating over the X matrix. - added script to create a synthetic dataset for testing - added a script to convert an existing CXG dataset to a sparse CXG dataset
This commit is contained in:
bmccandless
@@ -8,8 +8,11 @@ The organization of the TileDB structure is:
|
||||
├─ obs TileDB array containing cell (row) attributes, one attribute per
|
||||
│ dataframe column, shape (n_obs,)
|
||||
├─ var TileDB array containing gene (column) attributes, with one attribute per
|
||||
│ dataframe column, shape (n_obs,)
|
||||
│ dataframe column, shape (n_var,)
|
||||
├─ X Main count matrix as a 2D TileDB array, single unnamed numeric attribute
|
||||
├─ X_col_shift TilebDB Array used in column shift encoding, shape (n_var,), dtype = X.dtype.
|
||||
│ Single unnamed numeric attribute. If this array is sparse, and X_col_shift exists,
|
||||
│ then all values in the i'th column were subtracted by X_col_shift[i].
|
||||
├─ emb TileDB group, storing optional embeddings (group may be empty)
|
||||
│ └─ <name1> TileDB Array, single anon attribute, ND numeric array, shape (n_obs, N)
|
||||
└─ cxg_group_metadata Empty array used only to stash metadata about the overall object.
|
||||
@@ -70,6 +73,7 @@ import argparse
|
||||
import numpy as np
|
||||
from os.path import splitext, basename
|
||||
import json
|
||||
from scipy.stats import mode
|
||||
|
||||
from server.common.colors import convert_anndata_category_colors_to_cxg_category_colors
|
||||
from server.common.errors import ColorFormatException
|
||||
@@ -205,7 +209,7 @@ def write_cxg(
|
||||
log(1, "\t...embeddings created")
|
||||
|
||||
# X matrix
|
||||
save_X(container, adata, ctx, sparse_threshold)
|
||||
save_X(container, adata.X, ctx, sparse_threshold)
|
||||
log(1, "\t...X created")
|
||||
|
||||
|
||||
@@ -437,72 +441,146 @@ def create_X(X_name, shape, is_sparse):
|
||||
tiledb.DenseArray.create(X_name, schema)
|
||||
|
||||
|
||||
def check_sparse_X(adata, sparse_threshold):
|
||||
shape = adata.X.shape
|
||||
def evaluate_for_sparse_encoding(xdata, sparse_threshold):
|
||||
"""
|
||||
This function determines if the X matrix has a sparsity below the sparse_threshold.
|
||||
This function also returns the number of non-zeros encountered and number
|
||||
of elements evaluated. This function may return before evaluating the whole X matrix
|
||||
if it can be determined that X is not sparse enough.
|
||||
"""
|
||||
shape = xdata.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, :]
|
||||
a = xdata[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 (False, nnz, lim * shape[1])
|
||||
log(2, "\t...rows", lim, "of", shape[0], "nnz", nnz, "nnz percent %5.2f%%" % (100 * nnz / (lim * shape[1])))
|
||||
|
||||
return (nnz, shape[0] * shape[1])
|
||||
is_sparse = (100.0 * nnz / (shape[0] * shape[1])) < sparse_threshold
|
||||
return (is_sparse, nnz, shape[0] * shape[1])
|
||||
|
||||
|
||||
def save_X(container, adata, ctx, sparse_threshold):
|
||||
def evaluate_for_sparse_column_shift_encoding(xdata, sparse_threshold):
|
||||
"""Column shift encoding works by taking the most common value in each column, then
|
||||
subtracting that value from each element of the column. If each column mostly contains
|
||||
its most common value, then the resulting matrix can be very sparse.
|
||||
|
||||
This function determines if column shift encoding can be used to transform
|
||||
the X matrix into a sparse matrix with a sparsity below the sparse_threshold.
|
||||
If so, return the col_shift array that stores this encoding.
|
||||
This function also returns the number of non-zeros encountered and number
|
||||
of elements evaluated. This function may return before evaluating the whole X matrix
|
||||
if it can be determined that X cannot benefit from column shift encoding.
|
||||
"""
|
||||
shape = xdata.shape
|
||||
stride = max(1, 128_000_000 // shape[0])
|
||||
col_shift = np.zeros(shape[1])
|
||||
nnz = 0
|
||||
maxnnz = int(shape[0] * shape[1] * sparse_threshold / 100)
|
||||
for col in range(0, shape[1], stride):
|
||||
lim = min(col + stride, shape[1])
|
||||
a = xdata[:, col:lim]
|
||||
if type(a) is not np.ndarray:
|
||||
a = a.toarray()
|
||||
m = mode(a)
|
||||
col_shift[col:lim] = m.mode
|
||||
nnz += shape[0] * (lim - col) - np.sum(m.count)
|
||||
if nnz > maxnnz:
|
||||
return (None, nnz, shape[0] * lim)
|
||||
log(2, "\t...cols", lim, "of", shape[1], "nnz",
|
||||
nnz, "nnz percent %5.2f%%" % (100 * nnz / (lim * shape[0])))
|
||||
|
||||
is_sparse = (100.0 * nnz / (shape[0] * shape[1])) < sparse_threshold
|
||||
return (col_shift if is_sparse else None, nnz, shape[0] * shape[1])
|
||||
|
||||
|
||||
def save_X(container, xdata, ctx, sparse_threshold, expect_sparse=False):
|
||||
# Save X count matrix
|
||||
X_name = f"{container}/X"
|
||||
shape = adata.X.shape
|
||||
|
||||
shape = xdata.shape
|
||||
log(1, "\t...shape:", str(shape))
|
||||
|
||||
col_shift = None
|
||||
if sparse_threshold == 100:
|
||||
is_sparse = True
|
||||
elif sparse_threshold == 0:
|
||||
is_sparse = False
|
||||
else:
|
||||
nnz, nelem = check_sparse_X(adata, sparse_threshold)
|
||||
is_sparse, nnz, nelem = evaluate_for_sparse_encoding(xdata, 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)
|
||||
log(1, "\t...sparse=", is_sparse, "non-zeros percent (estimate): %6.2f" % percent)
|
||||
else:
|
||||
log(1, "\t...shape:", str(shape), "non-zeros:", nnz, "percent: %6.2f" % percent)
|
||||
log(1, "\t...sparse=", is_sparse, "non-zeros:", nnz, "percent: %6.2f" % percent)
|
||||
|
||||
is_sparse = percent < sparse_threshold
|
||||
if not is_sparse:
|
||||
col_shift, nnz, nelem = evaluate_for_sparse_column_shift_encoding(xdata, sparse_threshold)
|
||||
is_sparse = col_shift is not None
|
||||
percent = 100.0 * nnz / nelem
|
||||
if nelem != shape[0] * shape[1]:
|
||||
log(1, "\t...sparse=", is_sparse, "col shift non-zeros percent (estimate): %6.2f" % percent)
|
||||
else:
|
||||
log(1, "\t...sparse=", is_sparse, "col shift non-zeros:", nnz, "percent: %6.2f" % percent)
|
||||
|
||||
if expect_sparse is True and is_sparse is False:
|
||||
return False
|
||||
|
||||
create_X(X_name, shape, is_sparse)
|
||||
|
||||
stride = min(int(np.power(10, np.around(np.log10(1e9 / shape[1])))), 10_000)
|
||||
if is_sparse:
|
||||
log(1, "\t...output X as sparse matrix")
|
||||
if col_shift is not None:
|
||||
log(1, "\t...output X as sparse matrix with column shift encoding")
|
||||
X_col_shift_name = f"{container}/X_col_shift"
|
||||
filters = tiledb.FilterList([tiledb.ZstdFilter()])
|
||||
attrs = [tiledb.Attr(dtype=np.float32, filters=filters)]
|
||||
domain = tiledb.Domain(tiledb.Dim(domain=(0, shape[1] - 1), tile=min(shape[1], 5000), dtype=np.uint32))
|
||||
schema = tiledb.ArraySchema(domain=domain, attrs=attrs)
|
||||
tiledb.DenseArray.create(X_col_shift_name, schema)
|
||||
with tiledb.DenseArray(X_col_shift_name, mode="w", ctx=ctx) as X_col_shift:
|
||||
X_col_shift[:] = col_shift
|
||||
tiledb.consolidate(X_col_shift_name, ctx=ctx)
|
||||
else:
|
||||
log(1, "\t...output X as sparse matrix")
|
||||
|
||||
with tiledb.SparseArray(X_name, mode="w", ctx=ctx) as X:
|
||||
nnz = 0
|
||||
for row in range(0, shape[0], stride):
|
||||
lim = min(row + stride, shape[0])
|
||||
a = adata.X[row:lim, :]
|
||||
a = xdata[row:lim, :]
|
||||
if type(a) is not np.ndarray:
|
||||
a = a.toarray()
|
||||
if col_shift is not None:
|
||||
a = a - col_shift
|
||||
indices = np.nonzero(a)
|
||||
trow = indices[0] + row
|
||||
nnz += indices[0].shape[0]
|
||||
X[trow, indices[1]] = a[indices[0], indices[1]]
|
||||
log(2, "\t...rows", row, "to", lim)
|
||||
tiledb.consolidate(X_name, ctx=ctx)
|
||||
log(2, "\t...rows", lim, "of", shape[0], "nnz", nnz, "sparse", nnz / (lim * shape[1]))
|
||||
|
||||
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, :]
|
||||
a = xdata[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 hasattr(tiledb, "vacuum"):
|
||||
tiledb.vacuum(X_name)
|
||||
|
||||
return is_sparse
|
||||
|
||||
|
||||
def save_metadata(container, metadata_dict):
|
||||
|
||||
@@ -0,0 +1,62 @@
|
||||
"""
|
||||
Script to create a sparse dataset in CXG format based on an input dataset in CXG format.
|
||||
The input dataset is not modified.
|
||||
"""
|
||||
import os
|
||||
import shutil
|
||||
import tiledb
|
||||
import argparse
|
||||
import sys
|
||||
import server.converters.cxgtool as cxgtool
|
||||
|
||||
|
||||
def main():
|
||||
parser = argparse.ArgumentParser()
|
||||
parser.add_argument("input", help="input cxg directory")
|
||||
parser.add_argument("output", help="output cxg directory")
|
||||
parser.add_argument("--overwrite", action="store_true", help="replace output cxg directory")
|
||||
parser.add_argument("--verbose", "-v", action="count", default=0, help="verbose output")
|
||||
parser.add_argument(
|
||||
"--sparse-threshold",
|
||||
"-s",
|
||||
type=float,
|
||||
default=5.0, # default is 5% non-zero values
|
||||
help="The X array will be sparse if the percent of non-zeros falls below this value",
|
||||
)
|
||||
args = parser.parse_args()
|
||||
|
||||
if os.path.exists(args.output):
|
||||
print("output dir exists:", args.output)
|
||||
if args.overwrite:
|
||||
print("output dir removed:", args.output)
|
||||
shutil.rmtree(args.output)
|
||||
else:
|
||||
print("use the overwrite option to remove the output directory")
|
||||
sys.exit(1)
|
||||
|
||||
if not os.path.isdir(args.input):
|
||||
print("input is not a directory", args.input)
|
||||
sys.exit(1)
|
||||
|
||||
shutil.copytree(args.input, args.output,
|
||||
ignore=shutil.ignore_patterns("X", "X_col_shift"))
|
||||
|
||||
ctx = tiledb.Ctx(
|
||||
{
|
||||
"sm.num_reader_threads": 32,
|
||||
"sm.num_writer_threads": 32,
|
||||
"sm.consolidation.buffer_size": 1 * 1024 * 1024 * 1024,
|
||||
}
|
||||
)
|
||||
|
||||
with tiledb.DenseArray(os.path.join(args.input, "X"), mode="r", ctx=ctx) as X_in:
|
||||
is_sparse = cxgtool.save_X(args.output, X_in, ctx, args.sparse_threshold, expect_sparse=True)
|
||||
|
||||
if is_sparse is False:
|
||||
print("The array is not sparse, cleaning up, abort.")
|
||||
shutil.rmtree(args.output)
|
||||
sys.exit(1)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
Reference in New Issue
Block a user