novae 0.0.1__py3-none-any.whl

novae/__init__.py

@@ -0,0 +1,15 @@
+import importlib.metadata
+import logging
+
+__version__ = importlib.metadata.version("novae")
+
+from ._logging import configure_logger
+from .model import Novae
+from . import utils
+from . import data
+from . import monitor
+from . import plot
+from ._constants import settings
+
+log = logging.getLogger("novae")
+configure_logger(log)

novae/_constants.py

@@ -0,0 +1,87 @@
+import numpy as np
+
+
+class Keys:
+    # obs keys
+    LEAVES: str = "novae_leaves"
+    DOMAINS_PREFIX: str = "novae_domains_"
+    IS_VALID_OBS: str = "neighborhood_valid"
+    SLIDE_ID: str = "novae_sid"
+
+    # obsm keys
+    REPR: str = "novae_latent"
+    REPR_CORRECTED: str = "novae_latent_corrected"
+
+    # obsp keys
+    ADJ: str = "spatial_distances"
+    ADJ_LOCAL: str = "spatial_distances_local"
+    ADJ_PAIR: str = "spatial_distances_pair"
+
+    # var keys
+    VAR_MEAN: str = "mean"
+    VAR_STD: str = "std"
+    IS_KNOWN_GENE: str = "in_vocabulary"
+    HIGHLY_VARIABLE: str = "highly_variable"
+    USE_GENE: str = "novae_use_gene"
+
+    # layer keys
+    COUNTS_LAYER: str = "counts"
+
+    # misc keys
+    UNS_TISSUE: str = "novae_tissue"
+    ADATA_INDEX: str = "adata_index"
+    N_BATCHES: str = "n_batches"
+    NOVAE_VERSION: str = "novae_version"
+
+
+class Nums:
+    # training constants
+    EPS: float = 1e-8
+    MIN_DATASET_LENGTH: int = 50_000
+    MAX_DATASET_LENGTH_RATIO: float = 0.02
+    DEFAULT_SAMPLE_CELLS: int = 100_000
+    WARMUP_EPOCHS: int = 1
+
+    # distances constants and thresholds (in microns)
+    CELLS_CHARACTERISTIC_DISTANCE: int = 20  # characteristic distance between two cells, in microns
+    DELAUNAY_RADIUS_TH: int = 100
+    MEAN_DISTANCE_UPPER_TH_WARNING: float = 50
+    MEAN_DISTANCE_LOWER_TH_WARNING: float = 4
+
+    # genes constants
+    N_HVG_THRESHOLD: int = 500
+    MIN_GENES_FOR_HVG: int = 100
+    MIN_GENES: int = 20
+
+    # swav head constants
+    SWAV_EPSILON: float = 0.05
+    SINKHORN_ITERATIONS: int = 3
+    QUEUE_SIZE: int = 5
+    QUEUE_WEIGHT_THRESHOLD: float = 0.99
+
+    # misc nums
+    MEAN_NGH_TH_WARNING: float = 3.5
+    N_OBS_THRESHOLD: int = 2_000_000  # above this number, lazy loading is used
+
+    def disable_lazy_loading(self):
+        """Disable lazy loading of subgraphs in the NovaeDataset."""
+        Nums.N_OBS_THRESHOLD = np.inf
+
+    def enable_lazy_loading(self, n_obs_threshold: int = 0):
+        """Enable lazy loading of subgraphs in the NovaeDataset.
+
+        Args:
+            n_obs_threshold: Lazy loading is used above this number of cells in an AnnData object.
+        """
+        Nums.N_OBS_THRESHOLD = n_obs_threshold
+
+    @property
+    def warmup_epochs(self):
+        return Nums.WARMUP_EPOCHS
+
+    @warmup_epochs.setter
+    def warmup_epochs(self, value: int):
+        Nums.WARMUP_EPOCHS = value
+
+
+settings = Nums()

novae/_logging.py

@@ -0,0 +1,40 @@
+from __future__ import annotations
+
+import logging
+
+log = logging.getLogger(__name__)
+
+
+class ColorFormatter(logging.Formatter):
+    grey = "\x1b[38;20m"
+    blue = "\x1b[36;20m"
+    yellow = "\x1b[33;20m"
+    red = "\x1b[31;20m"
+    bold_red = "\x1b[31;1m"
+    reset = "\x1b[0m"
+
+    prefix = "[%(levelname)s] (%(name)s)"
+    suffix = "%(message)s"
+
+    FORMATS = {
+        logging.DEBUG: f"{grey}{prefix}{reset} {suffix}",
+        logging.INFO: f"{blue}{prefix}{reset} {suffix}",
+        logging.WARNING: f"{yellow}{prefix}{reset} {suffix}",
+        logging.ERROR: f"{red}{prefix}{reset} {suffix}",
+        logging.CRITICAL: f"{bold_red}{prefix}{reset} {suffix}",
+    }
+
+    def format(self, record):
+        log_fmt = self.FORMATS.get(record.levelno)
+        formatter = logging.Formatter(log_fmt)
+        return formatter.format(record)
+
+
+def configure_logger(log: logging.Logger):
+    log.setLevel(logging.INFO)
+
+    consoleHandler = logging.StreamHandler()
+    consoleHandler.setFormatter(ColorFormatter())
+
+    log.addHandler(consoleHandler)
+    log.propagate = False

novae/data/__init__.py

@@ -0,0 +1,3 @@
+from .convert import AnnDataTorch
+from .dataset import NovaeDataset
+from .datamodule import NovaeDatamodule

novae/data/convert.py

@@ -0,0 +1,143 @@
+from __future__ import annotations
+
+import numpy as np
+import torch
+from anndata import AnnData
+from sklearn.preprocessing import LabelEncoder
+from torch import Tensor
+
+from .._constants import Keys, Nums
+from ..module import CellEmbedder
+from ..utils import sparse_std
+
+
+class AnnDataTorch:
+    tensors: list[Tensor] | None
+    genes_indices_list: list[Tensor]
+
+    def __init__(self, adatas: list[AnnData], cell_embedder: CellEmbedder):
+        """Converting AnnData objects to PyTorch tensors.
+
+        Args:
+            adatas: A list of `AnnData` objects.
+            cell_embedder: A [novae.module.CellEmbedder][] object.
+        """
+        super().__init__()
+        self.adatas = adatas
+        self.cell_embedder = cell_embedder
+
+        self.genes_indices_list = [self._adata_to_genes_indices(adata) for adata in self.adatas]
+        self.tensors = None
+
+        self.means, self.stds, self.label_encoder = self._compute_means_stds()
+
+        # Tensors are loaded in memory for low numbers of cells
+        if sum(adata.n_obs for adata in self.adatas) < Nums.N_OBS_THRESHOLD:
+            self.tensors = [self.to_tensor(adata) for adata in self.adatas]
+
+    def _adata_to_genes_indices(self, adata: AnnData) -> Tensor:
+        return self.cell_embedder.genes_to_indices(adata.var_names[self._keep_var(adata)])[None, :]
+
+    def _keep_var(self, adata: AnnData) -> AnnData:
+        return adata.var[Keys.USE_GENE]
+
+    def _compute_means_stds(self) -> tuple[Tensor, Tensor, LabelEncoder]:
+        means, stds = {}, {}
+
+        for adata in self.adatas:
+            slide_ids = adata.obs[Keys.SLIDE_ID]
+            for slide_id in slide_ids.cat.categories:
+                adata_slide = adata[adata.obs[Keys.SLIDE_ID] == slide_id, self._keep_var(adata)]
+
+                mean = adata_slide.X.mean(0)
+                mean = mean.A1 if isinstance(mean, np.matrix) else mean
+                means[slide_id] = mean.astype(np.float32)
+
+                std = adata_slide.X.std(0) if isinstance(adata_slide.X, np.ndarray) else sparse_std(adata_slide.X, 0).A1
+                stds[slide_id] = std.astype(np.float32)
+
+        label_encoder = LabelEncoder()
+        label_encoder.fit(list(means.keys()))
+
+        means = [torch.tensor(means[slide_id]) for slide_id in label_encoder.classes_]
+        stds = [torch.tensor(stds[slide_id]) for slide_id in label_encoder.classes_]
+
+        return means, stds, label_encoder
+
+    def to_tensor(self, adata: AnnData, where_counts: Tensor | None = None) -> Tensor | tuple[Tensor, Tensor]:
+        """Get the normalized gene expressions of the cells in the dataset.
+        Only the genes of interest are kept (known genes and highly variable).
+
+        Args:
+            adata: An `AnnData` object.
+            where_counts: Where to keep expression as counts.
+
+        Returns:
+            A `Tensor` containing the normalized gene expresions.
+        """
+        adata = adata[:, self._keep_var(adata)]
+
+        if len(np.unique(adata.obs[Keys.SLIDE_ID])) == 1:
+            slide_id_index = self.label_encoder.transform([adata.obs.iloc[0][Keys.SLIDE_ID]])[0]
+            mean, std = self.means[slide_id_index], self.stds[slide_id_index]
+        else:
+            slide_id_indices = self.label_encoder.transform(adata.obs[Keys.SLIDE_ID])
+            mean = torch.stack([self.means[i] for i in slide_id_indices])  # TODO: avoid stack (only if not fast enough)
+            std = torch.stack([self.stds[i] for i in slide_id_indices])
+
+        X = adata.X if isinstance(adata.X, np.ndarray) else adata.X.toarray()
+        X = torch.tensor(X, dtype=torch.float32)
+        X = (X - mean) / (std + Nums.EPS)
+
+        if where_counts is None:
+            return X
+
+        counts = adata.layers[Keys.COUNTS_LAYER][:, where_counts]
+        counts = counts if isinstance(counts, np.ndarray) else counts.toarray()
+
+        return X[:, ~where_counts], torch.tensor(counts, dtype=torch.float32)
+
+    def __getitem__(self, item: tuple[int, slice]) -> tuple[Tensor, Tensor]:
+        """Get the expression values for a subset of cells (corresponding to a subgraph).
+
+        Args:
+            item: A `tuple` containing the index of the `AnnData` object and the indices of the cells in the neighborhoods.
+
+        Returns:
+            A `Tensor` of normalized gene expressions and a `Tensor` of gene indices.
+        """
+        adata_index, obs_indices = item
+
+        if self.tensors is not None:
+            return self.tensors[adata_index][obs_indices], self.genes_indices_list[adata_index]
+
+        adata = self.adatas[adata_index]
+        adata_view = adata[obs_indices]
+
+        return self.to_tensor(adata_view), self.genes_indices_list[adata_index]
+
+    def item_with_counts(
+        self, adata_index: int, obs_indices: slice, counts_ratio: float
+    ) -> tuple[Tensor, Tensor, Tensor]:
+        adata = self.adatas[adata_index]
+        adata_view = adata[obs_indices]
+
+        genes_indices = self.genes_indices_list[adata_index]
+
+        where_counts = _where_count(counts_ratio, len(genes_indices[0]))
+
+        X, counts = self.to_tensor(adata_view, where_counts)
+
+        return X, counts, genes_indices[:, ~where_counts], genes_indices[:, where_counts]
+
+
+def _where_count(counts_ratio: float, n_vars: int) -> Tensor:
+    n_vars_counts = int(counts_ratio * n_vars)
+    where_counts = torch.cat(
+        [
+            torch.ones(n_vars_counts, dtype=torch.bool),
+            torch.zeros(n_vars - n_vars_counts, dtype=torch.bool),
+        ]
+    )
+    where_counts = where_counts[torch.randperm(where_counts.size(0))]
+    return where_counts

novae/data/datamodule.py

@@ -0,0 +1,68 @@
+from __future__ import annotations
+
+import lightning as L
+from anndata import AnnData
+from torch_geometric.loader import DataLoader
+
+from ..module import CellEmbedder
+from . import NovaeDataset
+
+
+class NovaeDatamodule(L.LightningDataModule):
+    """
+    Datamodule used for training and inference. Small wrapper around the [novae.data.NovaeDataset][]
+    """
+
+    def __init__(
+        self,
+        adatas: list[AnnData],
+        cell_embedder: CellEmbedder,
+        batch_size: int,
+        n_hops_local: int,
+        n_hops_view: int,
+        num_workers: int = 0,
+        sample_cells: int | None = None,
+        counts_ratio: float | None = None,
+    ) -> None:
+        super().__init__()
+        self.dataset = NovaeDataset(
+            adatas,
+            cell_embedder=cell_embedder,
+            batch_size=batch_size,
+            n_hops_local=n_hops_local,
+            n_hops_view=n_hops_view,
+            sample_cells=sample_cells,
+            counts_ratio=counts_ratio,
+        )
+        self.batch_size = batch_size
+        self.num_workers = num_workers
+
+    def train_dataloader(self) -> DataLoader:
+        """Get a Pytorch dataloader for prediction.
+
+        Returns:
+            The training dataloader.
+        """
+        self.dataset.training = True
+        return DataLoader(
+            self.dataset,
+            batch_size=self.batch_size,
+            shuffle=False,
+            drop_last=True,
+            num_workers=self.num_workers,
+        )
+
+    def predict_dataloader(self) -> DataLoader:
+        """Get a Pytorch dataloader for prediction or inference.
+
+        Returns:
+            The prediction dataloader.
+        """
+        self.dataset.training = False
+        return DataLoader(
+            self.dataset,
+            batch_size=self.batch_size,
+            shuffle=False,
+            drop_last=False,
+            num_workers=self.num_workers,
+        )

novae/data/dataset.py

@@ -0,0 +1,254 @@
+from __future__ import annotations
+
+import numpy as np
+import pandas as pd
+import torch
+from anndata import AnnData
+from scipy.sparse import csr_matrix, lil_matrix
+from torch.utils.data import Dataset
+from torch_geometric.data import Data
+from torch_geometric.utils.convert import from_scipy_sparse_matrix
+
+from .. import utils
+from .._constants import Keys, Nums
+from ..module import CellEmbedder
+from . import AnnDataTorch
+
+
+class NovaeDataset(Dataset):
+    """
+    Dataset used for training and inference.
+
+    It extracts the the neighborhood of a cell, and convert it to PyTorch Geometric Data.
+
+    Attributes:
+        valid_indices (list[np.ndarray]): List containing, for each `adata`, an array that denotes the indices of the cells whose neighborhood is valid.
+        obs_ilocs (np.ndarray): An array of shape `(total_valid_indices, 2)`. The first column corresponds to the adata index, and the second column is the cell index for the corresponding adata.
+        shuffled_obs_ilocs (np.ndarray): same as obs_ilocs, but shuffled. Each batch will contain cells from the same slide.
+    """
+
+    valid_indices: list[np.ndarray]
+    obs_ilocs: np.ndarray
+    shuffled_obs_ilocs: np.ndarray
+
+    @utils.format_docs
+    def __init__(
+        self,
+        adatas: list[AnnData],
+        cell_embedder: CellEmbedder,
+        batch_size: int,
+        n_hops_local: int,
+        n_hops_view: int,
+        sample_cells: int | None = None,
+        counts_ratio: float | None = None,
+    ) -> None:
+        """_summary_
+
+        Args:
+            adatas: A list of `AnnData` objects.
+            cell_embedder: A [novae.module.CellEmbedder][] object.
+            batch_size: The model batch size.
+            {n_hops_local}
+            {n_hops_view}
+            sample_cells: If not None, the dataset if used to sample the subgraphs from precisely `sample_cells` cells.
+        """
+        super().__init__()
+        self.adatas = adatas
+        self.cell_embedder = cell_embedder
+        self.anndata_torch = AnnDataTorch(self.adatas, self.cell_embedder)
+
+        self.training = False
+
+        self.batch_size = batch_size
+        self.n_hops_local = n_hops_local
+        self.n_hops_view = n_hops_view
+        self.sample_cells = sample_cells
+        self.counts_ratio = counts_ratio
+
+        self.single_adata = len(self.adatas) == 1
+        self.single_slide_mode = self.single_adata and len(np.unique(self.adatas[0].obs[Keys.SLIDE_ID])) == 1
+
+        self._init_dataset()
+
+    def __repr__(self) -> str:
+        multi_slide_mode, multi_adata = not self.single_slide_mode, not self.single_adata
+        n_samples = sum(len(indices) for indices in self.valid_indices)
+        return f"{self.__class__.__name__} with {n_samples} samples ({multi_slide_mode=}, {multi_adata=})"
+
+    def _init_dataset(self):
+        for adata in self.adatas:
+            adjacency: csr_matrix = adata.obsp[Keys.ADJ]
+
+            if Keys.ADJ_LOCAL not in adata.obsp:
+                adata.obsp[Keys.ADJ_LOCAL] = _to_adjacency_local(adjacency, self.n_hops_local)
+            if Keys.ADJ_PAIR not in adata.obsp:
+                adata.obsp[Keys.ADJ_PAIR] = _to_adjacency_view(adjacency, self.n_hops_view)
+            if Keys.IS_VALID_OBS not in adata.obs:
+                adata.obs[Keys.IS_VALID_OBS] = adata.obsp[Keys.ADJ_PAIR].sum(1).A1 > 0
+
+        self.valid_indices = [utils.valid_indices(adata) for adata in self.adatas]
+
+        self.obs_ilocs = None
+        if self.single_adata:
+            self.obs_ilocs = np.array([(0, obs_index) for obs_index in self.valid_indices[0]])
+
+        self.slides_metadata: pd.DataFrame = pd.concat(
+            [
+                self._adata_slides_metadata(adata_index, obs_indices)
+                for adata_index, obs_indices in enumerate(self.valid_indices)
+            ],
+            axis=0,
+        )
+
+        self.shuffle_obs_ilocs()
+
+    def __len__(self) -> int:
+        if self.sample_cells is not None:
+            return min(self.sample_cells, len(self.shuffled_obs_ilocs))
+
+        if self.training:
+            n_obs = len(self.shuffled_obs_ilocs)
+            return min(n_obs, max(Nums.MIN_DATASET_LENGTH, int(n_obs * Nums.MAX_DATASET_LENGTH_RATIO)))
+
+        assert self.single_adata, "Multi-adata mode not supported for inference"
+
+        return len(self.obs_ilocs)
+
+    def __getitem__(self, index: int) -> dict[str, Data]:
+        """Gets a sample from the dataset, with one "main" graph and its corresponding "view" graph (only during training).
+
+        Args:
+            index: Index of the sample to retrieve.
+
+        Returns:
+            A dictionnary whose keys are names, and values are PyTorch Geometric `Data` objects. The `"view"` graph is only provided during training.
+        """
+        if self.training or self.sample_cells is not None:
+            adata_index, obs_index = self.shuffled_obs_ilocs[index]
+        else:
+            adata_index, obs_index = self.obs_ilocs[index]
+
+        data = self.to_pyg_data(adata_index, obs_index)
+
+        if not self.training or self.counts_ratio is not None:
+            return {"main": data}
+
+        adjacency_pair: csr_matrix = self.adatas[adata_index].obsp[Keys.ADJ_PAIR]
+        cell_view_index = np.random.choice(list(adjacency_pair[obs_index].indices), size=1)[0]
+
+        return {"main": data, "view": self.to_pyg_data(adata_index, cell_view_index)}
+
+    def to_pyg_data(self, adata_index: int, obs_index: int) -> Data:
+        """Create a PyTorch Geometric Data object for the input cell
+
+        Args:
+            adata_index: The index of the `AnnData` object
+            obs_index: The index of the input cell for the corresponding `AnnData` object
+
+        Returns:
+            A Data object
+        """
+        adata = self.adatas[adata_index]
+        adjacency_local: csr_matrix = adata.obsp[Keys.ADJ_LOCAL]
+        obs_indices = adjacency_local[obs_index].indices
+
+        adjacency: csr_matrix = adata.obsp[Keys.ADJ]
+        edge_index, edge_weight = from_scipy_sparse_matrix(adjacency[obs_indices][:, obs_indices])
+        edge_attr = edge_weight[:, None].to(torch.float32) / Nums.CELLS_CHARACTERISTIC_DISTANCE
+
+        if self.counts_ratio is None or not self.training:
+            x, genes_indices = self.anndata_torch[adata_index, obs_indices]
+            counts, counts_genes_indices = None, None
+        else:
+            x, counts, genes_indices, counts_genes_indices = self.anndata_torch.item_with_counts(
+                adata_index, obs_indices, self.counts_ratio
+            )
+
+        return Data(
+            x=x,
+            edge_index=edge_index,
+            edge_attr=edge_attr,
+            genes_indices=genes_indices,
+            slide_id=adata.obs[Keys.SLIDE_ID].iloc[0],
+            counts=counts,
+            counts_genes_indices=counts_genes_indices,
+        )
+
+    def shuffle_obs_ilocs(self):
+        """Shuffle the indices of the cells to be used in the dataset (for training only)."""
+        if self.single_slide_mode:
+            self.shuffled_obs_ilocs = self.obs_ilocs[np.random.permutation(len(self.obs_ilocs))]
+            return
+
+        adata_indices = np.empty((0, self.batch_size), dtype=int)
+        batched_obs_indices = np.empty((0, self.batch_size), dtype=int)
+
+        for uid in self.slides_metadata.index:
+            adata_index = self.slides_metadata.loc[uid, Keys.ADATA_INDEX]
+            adata = self.adatas[adata_index]
+            _obs_indices = np.where((adata.obs[Keys.SLIDE_ID] == uid) & adata.obs[Keys.IS_VALID_OBS])[0]
+            _obs_indices = np.random.permutation(_obs_indices)
+
+            n_elements = self.slides_metadata.loc[uid, Keys.N_BATCHES] * self.batch_size
+            if len(_obs_indices) >= n_elements:
+                _obs_indices = _obs_indices[:n_elements]
+            else:
+                _obs_indices = np.random.choice(_obs_indices, size=n_elements)
+
+            _obs_indices = _obs_indices.reshape((-1, self.batch_size))
+
+            adata_indices = np.concatenate([adata_indices, np.full_like(_obs_indices, adata_index)], axis=0)
+            batched_obs_indices = np.concatenate([batched_obs_indices, _obs_indices], axis=0)
+
+        permutation = np.random.permutation(len(batched_obs_indices))
+        adata_indices = adata_indices[permutation].flatten()
+        obs_indices = batched_obs_indices[permutation].flatten()
+
+        self.shuffled_obs_ilocs = np.stack([adata_indices, obs_indices], axis=1)
+
+    def _adata_slides_metadata(self, adata_index: int, obs_indices: list[int]) -> pd.DataFrame:
+        obs_counts: pd.Series = self.adatas[adata_index].obs.iloc[obs_indices][Keys.SLIDE_ID].value_counts()
+        slides_metadata = obs_counts.to_frame()
+        slides_metadata[Keys.ADATA_INDEX] = adata_index
+        slides_metadata[Keys.N_BATCHES] = (slides_metadata["count"] // self.batch_size).clip(1)
+        return slides_metadata
+
+
+def _to_adjacency_local(adjacency: csr_matrix, n_hops_local: int) -> csr_matrix:
+    """
+    Creates an adjancency matrix for which all nodes
+    at a distance inferior to `n_hops_local` are linked.
+    """
+    assert n_hops_local >= 1, f"n_hops_local must be greater than 0. Found {n_hops_local}."
+
+    adjacency_local: lil_matrix = adjacency.copy().tolil()
+    adjacency_local.setdiag(1)
+    for _ in range(n_hops_local - 1):
+        adjacency_local = adjacency_local @ adjacency
+    return adjacency_local.tocsr()
+
+
+def _to_adjacency_view(adjacency: csr_matrix, n_hops_view: int) -> csr_matrix:
+    """
+    Creates an adjacancy matrix for which all nodes separated by
+    precisely `n_hops_view` nodes are linked.
+    """
+    assert n_hops_view >= 1, f"n_hops_view must be greater than 0. Found {n_hops_view}."
+
+    if n_hops_view == 1:
+        adjacency_pair = adjacency.copy()
+        adjacency_pair.setdiag(0)
+        adjacency_pair.eliminate_zeros()
+        return adjacency_pair
+
+    adjacency_pair = adjacency.copy()
+    adjacency_pair.setdiag(1)
+    for i in range(n_hops_view - 1):
+        if i == n_hops_view - 2:
+            adjacency_previous = adjacency_pair.copy()
+        adjacency_pair = adjacency_pair @ adjacency
+    adjacency_pair = adjacency_pair.tolil()
+    adjacency_pair[adjacency_previous.nonzero()] = 0
+    adjacency_pair: csr_matrix = adjacency_pair.tocsr()
+    adjacency_pair.eliminate_zeros()
+    return adjacency_pair