Code Organization

The code is organized as follows:

cellarium/
└── ml/
    ├── "callbacks"        # Custom PyTorch Lightning callbacks
    ├── "core"             # Essential components
    │   ├── "CellariumModule"              # PyTorch Lightning Module for model, training step, and optimizer
    │   ├── "CellariumAnnDataDataModule"   # DataModule for multi-GPU DataLoader for AnnData objects
    │   └── "CellariumPipeline"            # Pipeline for data transformations and model inference
    ├── "data"             # Distributed AnnData Collection and multi-GPU Iterable Datasets
    ├── "lr_schedulers"    # Custom learning rate schedulers
    ├── "models"           # Cellarium ML models
    ├── "preprocessing"    # Pre-processing functions
    ├── "transforms"       # Data transformation modules
    ├── "utilities"        # Utility functions for various submodules
    └── "cli.py"           # Implements the "cellarium-ml" CLI. Models must be registered here

Installation

To install via pip:

pip install cellarium-ml

To install the developer version from source:

git clone https://github.com/cellarium-ai/cellarium-ml.git
cd cellarium-ml
make install  # runs pip install -e .[dev]

API Documentation and Tutorials

For detailed API documentation and tutorials, visit: Cellarium ML Documentation

For Developers

To run the tests:

make test-examples                   # runs single-device cli example tests
make test-dataloader                 # runs single-device dataloader related tests
TEST_DEVICES=2 make test-dataloader  # runs multi-device dataloader related test
make test                            # runs single-device (all other) tests
TEST_DEVICES=2 make test             # runs multi-device (all other) tests

To format the code automatically:

make format                # runs ruff formatter and fixes linter errors

To run the linters:

make lint                  # runs ruff linter and checks for formatter errors

To build the documentation:

make docs                  # builds the documentation at docs/build/html