Alternative on-disk dataset format #452

Hello everyone,

the current on-disk dataset format consists of one file per modality, plus extra files for pre-computed featurizer outputs. Because of this, the package has to load a lot of files at every execution, which also includes a lot of checks for structural integrity.

A lot of efforts for developing more efficient data structures have been made in the single cell space. While we are not working with single cell data here, the data structures are generally the same: variable x observation matrices. Variables can be genes for transcriptomics, or other features for other modalities. Observations are single cells in the single cell space, but would be samples /cell lines in the drevalpy case. This difference however is more a difference in the mental model, than a difference in data structure. So effectively, the single cell data structures can be used for the drevalpy data.

Now the most established data structure for single cell data is AnnData:

It has two axes: variables (usually genes) and observations (usually cells). var and obs are standard DataFrames that allow storing metadata about the genes and cells (e.g. tissue, condition, etc) respectively. X is the count matrix. There can be multiple layers, e.g. if one has raw counts, different types of normalized counts, etc. obsm is a space for storing multidimensional representations of the cells (e.g. PCA, UMAP, etc). The rest of the fields is not so relevant here.

A multimodal extension of this is MuData:
Unfortunately there is no nice image for MuData, so we will have to do more imagination. MuData objects are basically managed dictionaries of AnnData objects:

{
    expression: AnnData<obs x genes>,
    proteomics: AnnData<obs x proteins>,
    methylation: AnnData<obs x region>
}

So basically we see that we usually have the same set of observations (cells / samples) and different features on the var axis. This is again the exact same between single cell and bulk.

Now I mentioned that it is not just a dictionary, but a managed one. It is managed in the sense, that it can link the observations across modalities. The MuData object tracks the union of the obs identifies across modalities and allows DataFrame-style subsetting of all the modalities conveniently.

Now how would all of this be helpful in drevalpy?

First, it is important that we can store a drug-response matrix also in AnnData, with cell lines as obs and drugs as var, we can use X to store the primary response metric (e.g. IC50), and use layers to store alternative metrics (e.g. AUC). obs can be used to store all available information about cell lines, var can be used to store all information about drugs, also e.g. if they are combinations or not.

Fold splitting could then be done in one of two ways:

1. Real subsetting (adata_train = adata[train_mask, :] (LCO) or adata_train = adata[:, train_mask] (LDO))
2. Storing split information in the object: adata_train.obs['split'] = random.choice(...), or the same for var, respectively. One could also assign fold indices and then use all except the test index for training of course.

The above is a pure single-modality approach with the drug response data as the modality. However, if drug response would be a key of a multimodal MuData alongside the omics layers, we could perform the fold splitting in the exact same way, while keeping the omics layers perfectly aligned.

So basically we could build a single MuData object per dataset, which would then allow dropping most of the on-disk dataset handling code that is currently in the codebase.

Some cool side effects:

1. AnnData and MuData objects can be stored either as h5 (h5ad, h5mu) and zarr files. Both come with built-in compression, and zarr also supports lazy loading of data chunks.
2. AnnData and MuData objects can be concatenated with other objects of the same kind. So if we want to train on datasets A and B and evaluate on C, that becomes very convenient.