Concatenate datasets to a single array store¶

In the previous notebooks, we’ve seen how to incrementally create a collection of scRNA-seq datasets and train models on it.

Sometimes we want to concatenate all datasets into one big array to speed up ad-hoc queries for slices for arbitrary metadata (see this blog post). This is what CELLxGENE does to create Census: a number of .h5ad files are concatenated to give rise to a single tiledbsoma array store (CELLxGENE: scRNA-seq).

Note

This notebook is based on the tiledbsoma documentation.

import lamindb as ln
import pandas as pd
import tiledbsoma
import tiledbsoma.io
from functools import reduce

💡 connected lamindb: testuser1/test-scrna

ln.settings.transform.stem_uid = "oJN8WmVrxI8m"
ln.settings.transform.version = "1"
ln.track()

Query the collection of h5ad files that we’d like to convert into a single array.

collection = ln.Collection.filter(
    name="My versioned scRNA-seq collection", version="2"
).one()
collection.describe()

Prepare the array store¶

Prepare a path and a context for a new tiledbsoma.Experiment.

We will create our array store at the LaminDB instance root with folder name "scrna.tiledbsoma".

storage_settings = ln.setup.settings.storage
soma_path = (storage_settings.root / "scrna.tiledbsoma").as_posix()  # we could take any AWS S3 path, here

If our path is on AWS S3, we need to create a context with region information (exception: us-east-1). You can find more about tiledb configuration parameters in the tiledb documentation.

if storage_settings.type == "s3":  # if the storage location is on AWS S3
    storage_region = storage_settings.region
    ctx = tiledbsoma.SOMATileDBContext(tiledb_config={"vfs.s3.region": storage_region})
else:
    ctx = None

Prepare the AnnData objects¶

We need to prepare theAnnData objects in the collection to be concatenated into one tiledbsoma.Experiment. They need to have the same .var and .obs columns, .uns and .obsp should be removed.

adatas = [artifact.load() for artifact in collection.artifacts]

Compute the intersetion of all columns. All AnnData objects should have the same columns in their .obs, .var, .raw.var to be ingested into one tiledbsoma.Experiment.

obs_columns = reduce(pd.Index.intersection, [adata.obs.columns for adata in adatas])
var_columns = reduce(pd.Index.intersection, [adata.var.columns for adata in adatas])
var_raw_columns = reduce(pd.Index.intersection, [adata.raw.var.columns for adata in adatas])

Prepare the AnnData objects for concatenation. Prepare id fields, sanitize index names, intersect columns, drop slots. Here we have to drop .obsp, .uns and also columns from the dataframes that are not in the intersections obtained above, otherwise the ingestion will fail. We will need to provide obs and var names in tiledbsoma.io.register_anndatas, so we create these fileds (obs_id, var_id) from the dataframe indices.

for i, adata in enumerate(adatas):
    del adata.obsp
    del adata.uns
    
    adata.obs = adata.obs.filter(obs_columns)
    adata.obs["obs_id"] = adata.obs.index
    adata.obs["dataset"] = i
    adata.obs.index.name = None
    
    adata.var = adata.var.filter(var_columns)
    adata.var["var_id"] = adata.var.index
    adata.var.index.name = None
    
    drop_raw_var_columns = adata.raw.var.columns.difference(var_raw_columns)
    adata.raw.var.drop(columns=drop_raw_var_columns, inplace=True)
    adata.raw.var["var_id"] = adata.raw.var.index
    adata.raw.var.index.name = None

Create the array store¶

Register all the AnnData objects. Pass experiment_uri=None because tiledbsoma.Experiment doesn’t exist yet:

registration_mapping = tiledbsoma.io.register_anndatas(
    experiment_uri=None,
    adatas=adatas,
    measurement_name="RNA",
    obs_field_name="obs_id",
    var_field_name="var_id",
    append_obsm_varm=True
)

Ingest the AnnData objects sequentially, providing the context. This saves the AnnData objects in one array store.

for adata in adatas:
    tiledbsoma.io.from_anndata(
        experiment_uri=soma_path,
        anndata=adata,
        measurement_name="RNA",
        registration_mapping=registration_mapping,
        context=ctx
    )

Register the array store¶

Register the created tiledbsoma.Experiment store in lamindb:

soma_artifact = ln.Artifact(soma_path, description="My scRNA-seq SOMA Experiment").save()
soma_artifact.describe()

Query the array¶

Open and query the experiment. We can use the registered Artifact.

with soma_artifact.open() as soma_array:
    print(soma_array["obs"].read().concat().to_pandas())