mirror of
https://github.com/chanzuckerberg/cellxgene.git
synced 2026-09-17 05:47:58 +08:00
475 lines
17 KiB
Python
475 lines
17 KiB
Python
"""
|
|
This program converts an [AnnData H5AD](https://anndata.readthedocs.io/en/stable/)
|
|
into a cellxgene TileDB structure, aka a 'CXG'.
|
|
|
|
The organization of the TileDB structure is:
|
|
|
|
the.cxg TileDB Group
|
|
|-- obs TileDB array containing cell (row) attributes, one attribute per
|
|
| dataframe columm, shape (n_obs,)
|
|
|-- var TileDB array containing gene (column) attributes, with one attribute per
|
|
| dataframe column, shape (n_var,)
|
|
|-- X Main count matrix as a 2D TileDB array, single unnanmed numeric attribute
|
|
|-- 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.
|
|
...
|
|
|
|
All arrays are defined to have a uint32 domain, zero based. All X counds and embedding
|
|
coordinates are coerced to float32, which is ample precision for visualization purposes.
|
|
Dataframe (metadata) types are generally preserved, or where that is not possible,
|
|
converted to somemthing with equal representative value in the cellxgene application
|
|
(eg, categorical types are converted to string, bools to uint8, etc).
|
|
|
|
The following objects are also decorated with auxilliary metadata using TileDB
|
|
array metadata:
|
|
|
|
* cxg_group_metadata: minimally, will contain a 'cxg_version' field, which
|
|
is a semver string identifing the version number of the CXG layout.
|
|
It may also contain 'cxg_parameters', a JSON-encoded parameter list
|
|
describing CXG-wide dataset parameters.
|
|
|
|
* obs, var: both contain an optional 'cxg_schema' field that is a json string,
|
|
containing per-column (attribute) schema hinting. This is used where the TileDB
|
|
native typing information is insufficient to reconstruct useful information
|
|
such as categorical typing from Pandas DataFrames, and to communicate which column
|
|
is the preferred human-readable index for obs & var.
|
|
|
|
This file also embodies a number of empirically derived tiledb schema parameters,
|
|
including the global data layout, spatial tile size, and the like. The CXG is
|
|
self-describing in these areas, and the actual values (eg, tile size) are empirically
|
|
derived from benchmarking. They may change in the future.
|
|
|
|
---
|
|
|
|
TODO/ISSUES:
|
|
* add sub-command structure to argparse, for future sub-commands
|
|
* Possible future work: accept Loom files
|
|
|
|
"""
|
|
import re
|
|
import anndata
|
|
import tiledb
|
|
import argparse
|
|
import numpy as np
|
|
from os.path import splitext, basename
|
|
import json
|
|
|
|
|
|
# the CXG container version number. Must be a semver string.
|
|
CXG_VERSION = "0.1"
|
|
|
|
|
|
def log(level, *args):
|
|
global log_level
|
|
if log_level and level <= log_level:
|
|
print(*args)
|
|
|
|
|
|
def main():
|
|
parser = argparse.ArgumentParser()
|
|
parser.add_argument("h5ad", nargs="?", help="H5AD file name")
|
|
parser.add_argument(
|
|
"--backed", action="store_true", help="loaded in file backed mode. Will be slower, but use less memory."
|
|
)
|
|
parser.add_argument(
|
|
"--obs-names", help="Name of annotation to use for observations. If not specified, will use the obs index."
|
|
)
|
|
parser.add_argument(
|
|
"--var-names", help="Name of annotation to use for variables. If not specified, will use the var index."
|
|
)
|
|
parser.add_argument("--verbose", "-v", action="count", default=0, help="verbose output")
|
|
parser.add_argument("--title", help="Human readable dataset title. If omitted, will use filename")
|
|
parser.add_argument(
|
|
"--about",
|
|
metavar="<URL>",
|
|
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")
|
|
args = parser.parse_args()
|
|
|
|
global log_level
|
|
log_level = args.verbose
|
|
|
|
adata = anndata.read_h5ad(args.h5ad, backed="r" if args.backed else None)
|
|
log(1, f"{basename(args.h5ad)} loaded...")
|
|
|
|
basefname = splitext(basename(args.h5ad))[0]
|
|
out = args.out if args.out is not None else basefname
|
|
container = out if splitext(out)[1] == ".cxg" else out + ".cxg"
|
|
title = args.title if args.title is not None else basefname
|
|
|
|
write_cxg(adata, container, title, var_names=args.var_names, obs_names=args.obs_names, about=args.about)
|
|
|
|
log(1, "done")
|
|
|
|
|
|
def write_cxg(adata, container, title, var_names=None, obs_names=None, about=None):
|
|
if not adata.var.index.is_unique:
|
|
raise ValueError("Variable index is not unique - unable to convert.")
|
|
if not adata.obs.index.is_unique:
|
|
raise ValueError("Observation index is not unique - unable to convert.")
|
|
|
|
"""
|
|
TileDB bug TileDB-Inc/TileDB#1575 requires that we sanitize all column names
|
|
prior to saving. This can be reverted when the bug is fixed.
|
|
"""
|
|
log(0, "Warning: sanitizing all dataframe column names.")
|
|
clean_all_column_names(adata)
|
|
|
|
ctx = tiledb.Ctx(
|
|
{
|
|
"sm.num_reader_threads": 32,
|
|
"sm.num_writer_threads": 32,
|
|
"sm.consolidation.buffer_size": 1 * 1024 * 1024 * 1024,
|
|
}
|
|
)
|
|
|
|
tiledb.group_create(container, ctx=ctx)
|
|
log(1, f"\t...group created, with name {container}")
|
|
|
|
# dataset metadata
|
|
save_metadata(container, {"title": title, "about": about})
|
|
log(1, "\t...dataset metadata saved")
|
|
|
|
# var/gene dataframe
|
|
save_dataframe(container, "var", adata.var, var_names, ctx=ctx)
|
|
log(1, "\t...var dataframe created")
|
|
|
|
# obs/cell dataframe
|
|
save_dataframe(container, "obs", adata.obs, obs_names, ctx=ctx)
|
|
log(1, "\t...obs dataframe created")
|
|
|
|
# embeddings
|
|
e_container = f"{container}/emb"
|
|
tiledb.group_create(e_container, ctx=ctx)
|
|
save_embeddings(e_container, adata, ctx)
|
|
log(1, "\t...embeddings created")
|
|
|
|
# X matrix
|
|
save_X(container, adata, ctx)
|
|
log(1, "\t...X created")
|
|
|
|
|
|
"""
|
|
TODO: the code used to handle type inferencing should not be duplicated between
|
|
this tool and the server/common/utils code. When this tool is merged into
|
|
the cellxgene CLI, consolidate.
|
|
"""
|
|
|
|
|
|
def dtype_to_schema(dtype):
|
|
if dtype == np.float32:
|
|
return (np.float32, {})
|
|
elif dtype == np.int32:
|
|
return (np.int32, {})
|
|
elif dtype == np.bool_:
|
|
return (np.uint8, {type: "boolean"})
|
|
elif dtype == np.str:
|
|
return (np.unicode, {"type": "string"})
|
|
elif dtype == "category":
|
|
typ, hint = cxg_type(dtype.categories)
|
|
return (typ, {"type": "categorical", "categories": dtype.categories.tolist()})
|
|
else:
|
|
raise TypeError(f"Annotations of type {dtype} are unsupported.")
|
|
|
|
|
|
def _can_cast_to_float32(array):
|
|
if array.dtype.kind == "f":
|
|
# force downcast for all floats
|
|
return True
|
|
return False
|
|
|
|
|
|
def _can_cast_to_int32(array):
|
|
if array.dtype.kind in ["i", "u"]:
|
|
if np.can_cast(array.dtype, np.int32):
|
|
return True
|
|
ii32 = np.iinfo(np.int32)
|
|
if array.min() >= ii32.min and array.max() <= ii32.max:
|
|
return True
|
|
return False
|
|
|
|
|
|
def cxg_type(array):
|
|
try:
|
|
return dtype_to_schema(array.dtype)
|
|
except TypeError:
|
|
dtype = array.dtype
|
|
data_kind = dtype.kind
|
|
if _can_cast_to_float32(array):
|
|
return (np.float32, {})
|
|
elif _can_cast_to_int32(array):
|
|
return (np.int32, {})
|
|
elif data_kind == "O" and dtype == "object":
|
|
return (np.unicode, {"type": "string"})
|
|
else:
|
|
raise TypeError(f"Annotations of type {dtype} are unsupported.")
|
|
|
|
|
|
def cxg_dtype(array):
|
|
return cxg_type(array)[0]
|
|
|
|
|
|
def create_dataframe(name, df, ctx):
|
|
"""
|
|
Current access patterns are oriented toward reading very large slices of
|
|
the dataframe, one attribute at a time. Attribute data also tends to be
|
|
(often) repetitive (bools, categories, strings).
|
|
Given this, we use:
|
|
* a large tile size (1000)
|
|
* very aggressive compression levels
|
|
"""
|
|
filter = tiledb.FilterList(
|
|
[
|
|
# attempt aggressive compression as many of these dataframes are very repetitive
|
|
# strings, bools and other non-float data.
|
|
tiledb.ZstdFilter(level=22),
|
|
]
|
|
)
|
|
attrs = [tiledb.Attr(name=col, dtype=cxg_dtype(df[col]), filters=filter) for col in df]
|
|
domain = tiledb.Domain(tiledb.Dim(domain=(0, df.shape[0] - 1), tile=min(df.shape[0], 1000), dtype=np.uint32))
|
|
schema = tiledb.ArraySchema(
|
|
domain=domain, sparse=False, attrs=attrs, cell_order="row-major", tile_order="row-major"
|
|
)
|
|
tiledb.DenseArray.create(name, schema)
|
|
|
|
|
|
def create_unique_column_name(df_cols, col_name_prefix):
|
|
"""
|
|
given the columns of a dataframe, and a name prefix, return a column name which
|
|
does not exist in the dataframe, AND which is prefixed by `prefix`
|
|
|
|
The approach is to append a numeric suffix, starting at zero and increasing by
|
|
one, until an unused name is found (eg, prefix_0, prefix_1, ...).
|
|
"""
|
|
suffix = 0
|
|
while f"{col_name_prefix}{suffix}" in df_cols:
|
|
suffix += 1
|
|
return f"{col_name_prefix}{suffix}"
|
|
|
|
|
|
def alias_index_col(df, df_name, index_col_name):
|
|
"""
|
|
We rely in the existance of a unique, human-readable index for
|
|
any dataframe (eg, var is typically gene name, obs the cell name).
|
|
The user can specify these via the --obs-names and --var-names config.
|
|
If they are not specified, use the existing index to create them, giving
|
|
the resulting column a unique name (eg, "name").
|
|
|
|
In both cases, enforce that the result is unique, and communicate the
|
|
index column name via the 'index' field in the schema hints.
|
|
"""
|
|
if index_col_name is None:
|
|
if not df.index.is_unique:
|
|
raise KeyError(
|
|
f"Values in {df_name}.index must be unique. "
|
|
"Please prepare data to contain unique index values, or specify an "
|
|
"alternative with --{ax_name}-name."
|
|
)
|
|
index_col_name = create_unique_column_name(df.columns, "name_")
|
|
# turn the index into a normal column
|
|
df.rename_axis(index_col_name, inplace=True)
|
|
df.reset_index(inplace=True)
|
|
|
|
elif index_col_name in df.columns:
|
|
# User has specified alternative column for unique names, and it exists
|
|
if not df[index_col_name].is_unique:
|
|
raise KeyError(
|
|
f"Values in {df_name}.{index_col_name} must be unique. Please prepare data to contain unique values."
|
|
)
|
|
|
|
else:
|
|
raise KeyError(f"Annotation {index_col_name}, specified in --{df_name}-name, does not exist.")
|
|
|
|
return (df, index_col_name)
|
|
|
|
|
|
def save_dataframe(container, name, df, index_col_name, ctx):
|
|
A_name = f"{container}/{name}"
|
|
(df, index_col_name) = alias_index_col(df, name, index_col_name)
|
|
create_dataframe(A_name, df, ctx=ctx)
|
|
with tiledb.DenseArray(A_name, mode="w", ctx=ctx) as A:
|
|
value = {}
|
|
schema_hints = {}
|
|
for k, v in df.items():
|
|
dtype, hints = cxg_type(v)
|
|
value[k] = v.to_numpy(dtype=dtype)
|
|
if hints:
|
|
schema_hints.update({k: hints})
|
|
|
|
schema_hints.update({"index": index_col_name})
|
|
A[:] = value
|
|
A.meta["cxg_schema"] = json.dumps(schema_hints)
|
|
|
|
tiledb.consolidate(A_name, ctx=ctx)
|
|
|
|
|
|
def create_emb(e_name, emb):
|
|
"""
|
|
Embeddings are typically accessed with very large slices (or all of the embedding),
|
|
and do not benefit from overly aggressive compression due to their format. Given
|
|
this, we use:
|
|
* large tile size (1000)
|
|
* default compression level
|
|
"""
|
|
filters = tiledb.FilterList([tiledb.ZstdFilter()])
|
|
attrs = [tiledb.Attr(dtype=emb.dtype, filters=filters)]
|
|
dims = []
|
|
for d in range(emb.ndim):
|
|
shape = emb.shape
|
|
dims.append(tiledb.Dim("", domain=(0, shape[d] - 1), tile=min(shape[d], 1000), dtype=np.uint32))
|
|
domain = tiledb.Domain(*dims)
|
|
schema = tiledb.ArraySchema(
|
|
domain=domain, sparse=False, attrs=attrs, capacity=1_000_000, cell_order="row-major", tile_order="row-major"
|
|
)
|
|
tiledb.DenseArray.create(e_name, schema)
|
|
|
|
|
|
def is_valid_embedding(adata, name, arr):
|
|
""" return True if this layout data is a valid array for front-end presentation:
|
|
* ndarray, with shape (n_obs, >= 2), dtype float/int/uint
|
|
* contains only finite values
|
|
* follows ScanPy embedding naming conventions
|
|
"""
|
|
is_valid = type(name) == str and name.startswith("X_") and len(name) > 2
|
|
is_valid = is_valid and type(arr) == np.ndarray and arr.dtype.kind in "fiu"
|
|
is_valid = is_valid and arr.shape[0] == adata.n_obs and arr.shape[1] >= 2
|
|
is_valid = is_valid and np.all(np.isfinite(arr))
|
|
return is_valid
|
|
|
|
|
|
def save_embeddings(container, adata, ctx):
|
|
for (name, value) in adata.obsm.items():
|
|
if is_valid_embedding(adata, name, value):
|
|
e_name = f"{container}/{name[2:]}"
|
|
create_emb(e_name, value)
|
|
with tiledb.DenseArray(e_name, mode="w", ctx=ctx) as A:
|
|
A[:] = value
|
|
tiledb.consolidate(e_name, ctx=ctx)
|
|
log(1, f"\t\t...{name} embedding created")
|
|
|
|
|
|
def create_X(X_name, shape):
|
|
"""
|
|
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
|
|
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 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),
|
|
)
|
|
schema = tiledb.ArraySchema(
|
|
domain=domain, sparse=False, attrs=attrs, cell_order="row-major", tile_order="col-major"
|
|
)
|
|
tiledb.DenseArray.create(X_name, schema)
|
|
|
|
|
|
def save_X(container, adata, ctx):
|
|
# Save X count matrix
|
|
X_name = f"{container}/X"
|
|
shape = adata.X.shape
|
|
create_X(X_name, shape)
|
|
|
|
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)
|
|
|
|
tiledb.consolidate(X_name, ctx=ctx)
|
|
|
|
|
|
def save_metadata(container, metadata):
|
|
"""
|
|
Save all dataset-wide metadata. This includes:
|
|
* CXG version
|
|
* dataset metadata, such as title and about link.
|
|
|
|
Longer term, tiledb will have support for metadata on groups. Until
|
|
such feature exists, create an empty array and annotate that array.
|
|
|
|
https://github.com/TileDB-Inc/TileDB-Py/issues/254
|
|
"""
|
|
a_name = f"{container}/cxg_group_metadata"
|
|
with tiledb.from_numpy(a_name, np.zeros((1,))) as A:
|
|
pass
|
|
with tiledb.DenseArray(a_name, mode="w") as A:
|
|
A.meta["cxg_version"] = CXG_VERSION
|
|
A.meta["cxg_properties"] = json.dumps(metadata)
|
|
|
|
|
|
def sanitize_keys(keys):
|
|
"""
|
|
We need names to be safe to use as attribute names in tiledb. See:
|
|
TileDB-Inc/TileDB#1575
|
|
TileDB-Inc/TileDB-Py#294
|
|
This can be entirely removed once they add proper escaping.
|
|
|
|
Args: list of keys
|
|
Returns: dict of {old_key: new_key, ...}
|
|
|
|
Returned new keys will be both safe and unique.
|
|
|
|
Masking out [~/.] and anything outside the ASCII range.
|
|
"""
|
|
p = re.compile(r"[^ -\.0-\[\]-\}]")
|
|
clean_keys = {k: p.sub("_", k) for k in keys}
|
|
|
|
used_keys = set()
|
|
clean_unique_keys = {}
|
|
for k, v in clean_keys.items():
|
|
if v not in used_keys:
|
|
used_keys.add(v)
|
|
clean_unique_keys[k] = v
|
|
continue
|
|
|
|
# else, needs deduping.
|
|
counter = 1
|
|
while True:
|
|
candidate_name = v + "-" + str(counter)
|
|
if candidate_name not in used_keys:
|
|
used_keys.add(candidate_name)
|
|
clean_unique_keys[k] = candidate_name
|
|
break
|
|
counter += 1
|
|
|
|
for k, v, in clean_unique_keys.items():
|
|
if k != v:
|
|
log(1, f"Renaming {k} to {v}")
|
|
return clean_unique_keys
|
|
|
|
|
|
def sanitize_df(df):
|
|
df.rename(columns=sanitize_keys(df.keys().tolist()), inplace=True)
|
|
|
|
|
|
def sanitize_mapping(mapping):
|
|
clean_keys = sanitize_keys([k for k in mapping.keys()])
|
|
for old_key, new_key in clean_keys.items():
|
|
if old_key != new_key:
|
|
mapping[new_key] = mapping[old_key]
|
|
del mapping[old_key]
|
|
|
|
|
|
def clean_all_column_names(adata):
|
|
sanitize_df(adata.obs)
|
|
sanitize_df(adata.var)
|
|
sanitize_mapping(adata.obsm)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|