CellFilter transform to eliminate cells during training

[sjfleming]

There is a possibility that after applying some arbitrary gene filtering, scVI would encounter cells with only a few counts... or even zero counts! The upfront cell filtering we do in Nexus looks for cells with some minimal number of UMIs or nonzero genes. But there is no guarantee this doesn't change after Filtering to just 4k or 8k genes. If we are encountering cells with extremely low UMIs during training, we should probably skip them.

OLD:
This PR introduces min_count_per_cell_threshold to allow us to skip such cells during training (on the fly).

EDIT:
This PR introduces the CellFilter Transform which can eliminate cells from a minibatch on the fly. It can be used with any model. It cannot be used with the contrastive learning model (due to shape expectations).

[sjfleming]

Hmmm... you know... this could alternatively be implemented more broadly than in scVI by making it a cellarium-ml Transform... like a CellFilter

[sjfleming]

Actually... eliminating user cells is a potential problem here in this PR as written. If a user approaches a trained model with 10k cells, this model could return < 10k embeddings, which is not good.

[sjfleming]

It might be easier / conceptually clearer at inference time to modify a Transform than to reach into the saved scVI model and start changing attributes.

[fedorgrab]

Hmmm... you know... this could alternatively be implemented more broadly than in scVI by making it a cellarium-ml Transform... like a CellFilter

I like that idea. Filtering on the model side seems a little more disorganized. Keeping the filtering on the transform side preserves the original separation of responsibilities: transforms handle data preprocessing, while the model does what a model is supposed to do.

[sjfleming]

Okay so the intention here is also to allow for different CellFilter behavior automatically at prediction time. If prediction is run via a CellariumModule.forward (which calls CellariumPipeline.predict), then a marker is set in batch which tells CellFilter (and any other method that may care eventually) that we are in predict mode. In that case, CellFilter is a no-op. The idea is that when we ask for output for a set of cells during prediction, we want them all.

Alternatively, one could reach into the transform's attributes and change the filters to zeros to get a no-op, but this is annoying (and specific to CellFilter). The different behavior on predict seems to make better sense (if it works well).