Flexible featurizer-predictor combinations and registries

[nictru]

I will describe here how I implemented the flexible combinations in my effort, so we can build a consensus on how this can be implemented in drevalpy. The terms 'featurizer' (what they are called in drevalpy) and 'encoder' (what I call them) are used interchangeably. Similar for drug/perturbation.

Registries

There are three registries: One for omics featurizers, one for perturbation/drug featurizers and one for predictors. Each featurizer/predictor registers itself to the corresponding registry using a decorator:

@register_omics_encoder("pca", description="...")
class PCA(OmicsEncoder): ...

@register_predictor("randomForest", description="...")
class RandomForestPredictor(SklearnTabularPredictor): ...

For components that are based on literature models, it is possible to store additional metadata, such as a publication DOI or a github link in the registry. This allows getting a comprehensive overview of the available components via a CLI command or python API.

Extending the registry

During development, having to modify the dreval source code for adding new models is annoying (and users might modify parts of the package that are not relevant for adding models). This is even more problematic for evaluating new components in the nextflow pipeline.

Registries can help here. By importing the decorator function from the package, users can extend the registry without touching the drevalpy source code. This can look like following:

import anndata as ad
import numpy as np
from sklearn.linear_model import Ridge

from drevalpy.models.encoders.omics.base import OmicsEncoder
from drevalpy.models.predictors.sklearn_tabular import SklearnTabularPredictor
from drevalpy.models.registry import register_omics_encoder, register_predictor


@register_omics_encoder(
    "logGeneExpression",
    description="Log1p of one obsm matrix.",
    category="native",
)
class LogGeneExpression(OmicsEncoder):
    def __init__(self, *, obsm_key: str = "gene_expression") -> None:
        self._obsm_key = obsm_key
        self._n_features = 0

    def _fit_impl(self, adata: ad.AnnData, *, row_indices=None) -> None:
        self._n_features = adata.obsm[self._obsm_key].shape[1]

    def transform(self, adata: ad.AnnData) -> np.ndarray:
        x = np.asarray(adata.obsm[self._obsm_key], dtype=np.float32)
        return np.log1p(np.clip(x, 0.0, None))

    @property
    def output_dim(self) -> int:
        return self._n_features


@register_predictor(
    "myRidge",
    description="Ridge on concat(omics, perturbation).",
    category="general_purpose",
)
class MyRidgePredictor(SklearnTabularPredictor):
    def _make_estimator(self) -> Ridge:
        return Ridge(alpha=float(self._h.get("alpha", 1.0)))

And then it can be used as follows:

import my_extensions  # registers logGeneExpression + myRidge

from drevalpy.models import EncoderConfig, ModelConfig, PredictorConfig

config = ModelConfig(
    omics_encoder=EncoderConfig("logGeneExpression", registry="omics"),
    perturbation_encoder=EncoderConfig("oneHot", registry="perturbation"),
    predictor=PredictorConfig("myRidge", hyperparameters={"alpha": 0.5}),
)

model = config.create_model()
model.fit(adata, strategy, fold_index=0)

When used via CLI (or via the nextflow pipeline) users can provide a directory containing such files, which are then automatically registered and usable.

Flexible combinations

Now as every component can be accessed via a string in its registry, it is possible to construct a model using a structured string. I used the pattern <omics encoder>:<drug_encoder>:<predictor>. We can then dynamically construct the equivalent of a DRPmodel dynamically. Each component comes with a defined hyperparameter space, and the hyperparameter space of the triplet is the union of the per-component hyperparameters, which can be tuned in combination.

Model zoo

Certain predictors (e.g. PharmaFormer) might require a certain pair of featurizers so we can really call it PharmaFormer. To allow giving certain combinations a name, I created a so-called 'model zoo' which stores certain combinations under a friendly name. It is also configure hyperparameter spaces and default hyperparameters like this.

It is also possible to provide additional model zoo configurations via a zoo directory to the CLI or the nextflow pipeline. It is also possible to reference self-registered components in the self-made zoo configurations.

Featurizer type checking

Most predictors get dense representations from their featurizers, but some also require graphs or sequence representations. The registry decorator requires a featurizer type for each featurizer, and an omics/perturbation featurizer type for each predictor. During construction of the combined model, we can thus ensure that a valid combination was selected before running into an actual type mismatch.

---

Let me know what you think about this

[JudithBernett]

At first glance, I really like it, but I think we need to think about this thoroughly, this is probably where your code diverged the most from ours. But I also think that requiring the user to fork our repo and modify the source code isn't super pretty compared to this solution! So I think it sounds really nice and it would also make integrating SynVerse much easier :)

[nictru]

Yes I think so too - I still need to look at the synverse code to understand how they do it

[nictru]

I had a look at the synverse code now, and my conclusion is that what they do is complementary to what I propose here. What I propose here is kind of a restructure of the featurizer/predictor implementations to make it more convenient to add new ones and build combinations of them.

Disclaimer:
For my analysis of synverse I ignored all the features that are exclusive to synergy prediction, as they are not very relevant for my work in this space. It could still make sense to add them to the package, but I don't have the resources for doing it.

From what I can see, these are the main features that SynVerse adds beyond what I already proposed:

1. Preprocessing policies
2. Sweep configurations
3. Ablation support

---

In the following, I explain each of them in more detail:

Preprocessing policies

In both drevalpy and my extended implementation, we have drug featurizers, cell line featurizers and predictors. SynVerse adds an additional layer between featurizers and predictors, called preprocessors. These can be any type of normalization, dimensionality reduction, compression, or similar.

Sweep configurations

While my suggested improvement allows building new combinations easily, it requires explicitly adding each desired combination to the run configuration explicitly. SynVerse in contrast has a layer of abstraction on top of that, which allows to select the desired featurizers, predictors, preprocessors, etc. and it will evaluate all of the possible concrete model configurations that arise from the requrested components.

Ablation support

They provide a simple interface for configuring ablations, e.g. running a training setup with both the proper features, and shuffled versions.

---

From my perspective, these features are all compatible with and even benefit from implementing my suggested registry-based approach for model construction. The sweep configuration would be a simple next step, which we could implement by accepting either an explicit list of model configurations or a sweep configuration. We can then create a dedicated CLI endpoint that allows generating the explicit list from the sweep configuration. Everything else could stay the same.

The preprocessing policies could potentially be implemented via featurizer hyperparameters. I think this could potentially be more elegant than adding another type of component between featurizers and predictors. But I'm open to discuss this.

The ablation support could probably be implemented fairly easily by modfiying the featurizer base classes to provide a randomization option, which can then be toggled via run configuration.

---

In summary, I think that the synverse features do not conflict with my suggested changes, so I will start working on implementing the registry-based architecture now.

[PascalIversen]

Hey, many thanks for the thoughts and the write-up! Starting to put some thought into this now. Something I do not understand yet: Many encoders are trained end-2-end (e.g. DrugGNN includes a drug graph encoder which is in the same comp. Graph as the predictor model). Is this then not an encoder in your vocab?

For me, this distinction makes sense:
1.)featurizer: fixed. can be calculated prior to the training
2.)encoder: part of the model and is optimized alongside it. can't be pre-calculated

As I understand it, currently this change would address 1, but not 2. (I guess that is fully okay for now, just wanted to spell out my thoughts)

The featurizer type check is a bit underpowered in the sense that a graph can have node or edge features or not, or one model might take in one cell-cell similarity global graph, and another many different drug graphs. Not sure how confusing that can get

[nictru]

In my implementation (but using your vocab concerning featurizers/encoders), drugGNN is implemented as follows:

1. Featurizer that takes SMILES and produces a PyG graph
2. The actual graph encoder that is optimized alongside the model is part of the model (or predictor in my terminology)

---

About the details on the graph compatibility checks:
Yes it's underpowered in my current (and also in my implementation), and doing it in a really robust way is probably quite hard. So far I have three graph featurizers: the drugGraph featurizer (what I mentioned for drugGNN), a ppiGraph and a pathwayGraph (both give a single graph per cell line).

I think there are two ways of handling this:

1. Implement it with the underpowered compatibility check and potentially strengthen later
2. Do not implement graph featurizers at all, but instead use an identity featurizer (that just passes the raw omics) and make the graph construction part of the model, could still be a reusable Mixin or something like that

[JudithBernett]

I think we can proceed with this implementation and then tackle these graph compatibility problems as we go, if someone has a great new idea for a graph representation of a drug/cell line that would not be compatible.