scdataloader 2.0.2__tar.gz → 2.0.4__tar.gz

{scdataloader-2.0.2 → scdataloader-2.0.4}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: scdataloader
-Version: 2.0.2
+Version: 2.0.4
 Summary: a dataloader for single cell data in lamindb
 Project-URL: repository, https://github.com/jkobject/scDataLoader
 Author-email: jkobject <jkobject@gmail.com>

{scdataloader-2.0.2 → scdataloader-2.0.4}/pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "scdataloader"
-version = "2.0.2"
+version = "2.0.4"
 description = "a dataloader for single cell data in lamindb"
 authors = [
     {name = "jkobject", email = "jkobject@gmail.com"}
@@ -15,6 +15,7 @@ dependencies = [
     "cellxgene-census>=0.1.0",
     "torch>=2.2.0",
     "pytorch-lightning>=2.3.0",
+    "lightning>=2.3.0",
     "anndata>=0.9.0",
     "zarr>=2.10.0",
     "matplotlib>=3.5.0",
@@ -27,8 +28,6 @@ dependencies = [
     "django>=4.0.0",
     "scikit-misc>=0.5.0",
     "jupytext>=1.16.0",
-    "lightning>=2.3.0",
-    "pytorch-lightning>=2.3.0",
 ]
 
 [project.optional-dependencies]

{scdataloader-2.0.2 → scdataloader-2.0.4}/scdataloader/collator.py

@@ -27,37 +27,60 @@ class Collator:
         genedf: Optional[pd.DataFrame] = None,
     ):
         """
-        This class is responsible for collating data for the scPRINT model. It handles the
-        organization and preparation of gene expression data from different organisms,
-        allowing for various configurations such as maximum gene list length, normalization,
-        and selection method for gene expression.
+        Collator for preparing gene expression data batches for the scPRINT model.
 
-        This Collator should work with scVI's dataloader as well!
+        This class handles the organization and preparation of gene expression data from
+        different organisms, allowing for various configurations such as maximum gene list
+        length, normalization, binning, and gene selection strategies.
+
+        Compatible with scVI's dataloader and other PyTorch data loading pipelines.
 
         Args:
-            organisms (list): List of organisms to be considered for gene expression data.
-                it will drop any other organism it sees (might lead to batches of different sizes!)
-            how (flag, optional): Method for selecting gene expression. Defaults to "most expr".
-                one of ["most expr", "random expr", "all", "some"]:
-                "most expr": selects the max_len most expressed genes,
-                if less genes are expressed, will sample random unexpressed genes,
-                "random expr": uses a random set of max_len expressed genes.
-                if less genes are expressed, will sample random unexpressed genes
-                "all": uses all genes
-                "some": uses only the genes provided through the genelist param
-            org_to_id (dict): Dictionary mapping organisms to their respective IDs.
-            valid_genes (list, optional): List of genes from the datasets, to be considered. Defaults to [].
-                it will drop any other genes from the input expression data (usefull when your model only works on some genes)
-            max_len (int, optional): Total number of genes to use (for random expr and most expr). Defaults to 2000.
-            n_bins (int, optional): Number of bins for binning the data. Defaults to 0. meaning, no binning of expression.
-            add_zero_genes (int, optional): Number of additional unexpressed genes to add to the input data. Defaults to 0.
-            logp1 (bool, optional): If True, logp1 normalization is applied. Defaults to False.
-            norm_to (float, optional): Rescaling value of the normalization to be applied. Defaults to None.
-            organism_name (str, optional): Name of the organism ontology term id. Defaults to "organism_ontology_term_id".
-            tp_name (str, optional): Name of the heat diff. Defaults to None.
-            class_names (list, optional): List of other classes to be considered. Defaults to [].
-            genelist (list, optional): List of genes to be considered. Defaults to [].
-                If [] all genes will be considered
+            organisms (List[str]): List of organism ontology term IDs to include.
+                Samples from other organisms will be dropped (may lead to variable batch sizes).
+            how (str, optional): Gene selection strategy. Defaults to "all".
+                - "most expr": Select the `max_len` most expressed genes. If fewer genes
+                  are expressed, randomly sample unexpressed genes to fill.
+                - "random expr": Randomly select `max_len` expressed genes. If fewer genes
+                  are expressed, randomly sample unexpressed genes to fill.
+                - "all": Use all genes without filtering.
+                - "some": Use only genes specified in the `genelist` parameter.
+            org_to_id (dict[str, int], optional): Mapping from organism names to integer IDs.
+                If None, organism names are used directly. Defaults to None.
+            valid_genes (List[str], optional): List of gene names to consider from input data.
+                Genes not in this list will be dropped. Useful when the model only supports
+                specific genes. Defaults to None (use all genes).
+            max_len (int, optional): Maximum number of genes to include when using "most expr"
+                or "random expr" selection methods. Defaults to 2000.
+            add_zero_genes (int, optional): Number of additional unexpressed genes to include
+                in the output. Only applies when `how` is "most expr" or "random expr".
+                Defaults to 0.
+            logp1 (bool, optional): Apply log2(1 + x) transformation to expression values.
+                Applied after normalization if both are enabled. Defaults to False.
+            norm_to (float, optional): Target sum for count normalization. Expression values
+                are scaled so that total counts equal this value. Defaults to None (no normalization).
+            n_bins (int, optional): Number of bins for expression value binning. If 0, no
+                binning is applied. Binning uses quantile-based discretization. Defaults to 0.
+            tp_name (str, optional): Column name in batch data for time point or heat diffusion
+                values. If None, time point values default to 0. Defaults to None.
+            organism_name (str, optional): Column name in batch data for organism ontology
+                term ID. Defaults to "organism_ontology_term_id".
+            class_names (List[str], optional): List of additional metadata column names to
+                include in the output. Defaults to [].
+            genelist (List[str], optional): List of specific genes to use when `how="some"`.
+                Required if `how="some"`. Defaults to [].
+            genedf (pd.DataFrame, optional): DataFrame containing gene information indexed by
+                gene name with an 'organism' column. If None, loaded automatically using
+                `load_genes()`. Defaults to None.
+
+        Attributes:
+            organism_ids (set): Set of organism IDs being processed.
+            start_idx (dict): Mapping from organism ID to starting gene index in the model.
+            accepted_genes (dict): Boolean masks for valid genes per organism.
+            to_subset (dict): Boolean masks for genelist filtering per organism.
+
+        Raises:
+            AssertionError: If `how="some"` but `genelist` is empty.
         """
         self.organisms = organisms
         self.max_len = max_len
@@ -77,6 +100,19 @@ class Collator:
         self._setup(genedf, org_to_id, valid_genes, genelist)
 
     def _setup(self, genedf=None, org_to_id=None, valid_genes=[], genelist=[]):
+        """
+        Initialize gene mappings and indices for each organism.
+
+        Sets up internal data structures for gene filtering, organism-specific
+        gene indices, and gene subsetting based on the provided configuration.
+
+        Args:
+            genedf (pd.DataFrame, optional): Gene information DataFrame. If None,
+                loaded via `load_genes()`. Defaults to None.
+            org_to_id (dict, optional): Organism name to ID mapping. Defaults to None.
+            valid_genes (List[str], optional): Genes to accept from input. Defaults to [].
+            genelist (List[str], optional): Genes to subset to when `how="some"`. Defaults to [].
+        """
         if genedf is None:
             genedf = load_genes(self.organisms)
             self.organism_ids = (
@@ -108,18 +144,45 @@ class Collator:
 
     def __call__(self, batch) -> dict[str, Tensor]:
         """
-        __call__ applies the collator to a minibatch of data
+        Collate a minibatch of gene expression data.
+
+        Processes a list of sample dictionaries, applying gene selection, normalization,
+        log transformation, and binning as configured. Filters out samples from organisms
+        not in the configured organism list.
 
         Args:
-            batch (List[dict[str: array]]): List of dicts of arrays containing gene expression data.
-                the first list is for the different samples, the second list is for the different elements with
-                elem["X"]: gene expression
-                elem["organism_name"]: organism ontology term id
-                elem["tp_name"]: heat diff
-                elem["class_names.."]: other classes
+            batch (List[dict]): List of sample dictionaries, each containing:
+                - "X" (array): Gene expression values.
+                - organism_name (any): Organism identifier (column name set by `organism_name`).
+                - tp_name (float, optional): Time point value (column name set by `tp_name`).
+                - class_names... (any, optional): Additional class labels.
+                - "_storage_idx" (int, optional): Dataset storage index.
+                - "is_meta" (int, optional): Metadata flag.
+                - "knn_cells" (array, optional): KNN neighbor expression data.
+                - "knn_cells_info" (array, optional): KNN neighbor metadata.
 
         Returns:
-            List[Tensor]: List of tensors containing the collated data.
+            dict[str, Tensor]: Dictionary containing collated tensors:
+                - "x" (Tensor): Gene expression matrix of shape (batch_size, n_genes).
+                  Values may be raw counts, normalized, log-transformed, or binned
+                  depending on configuration.
+                - "genes" (Tensor): Gene indices of shape (batch_size, n_genes) as int32.
+                  Indices correspond to positions in the model's gene vocabulary.
+                - "class" (Tensor): Class labels of shape (batch_size, n_classes) as int32.
+                - "tp" (Tensor): Time point values of shape (batch_size,).
+                - "depth" (Tensor): Total counts per cell of shape (batch_size,).
+                - "is_meta" (Tensor, optional): Metadata flags as int32. Present if input
+                  contains "is_meta".
+                - "knn_cells" (Tensor, optional): KNN expression data. Present if input
+                  contains "knn_cells".
+                - "knn_cells_info" (Tensor, optional): KNN metadata. Present if input
+                  contains "knn_cells_info".
+                - "dataset" (Tensor, optional): Dataset indices as int64. Present if input
+                  contains "_storage_idx".
+
+        Note:
+            Batch size in output may be smaller than input if some samples are filtered
+            out due to organism mismatch.
         """
         # do count selection
         # get the unseen info and don't add any unseen

{scdataloader-2.0.2 → scdataloader-2.0.4}/scdataloader/data.py

@@ -22,30 +22,66 @@ from .mapped import MappedCollection, _Connect
 @dataclass
 class Dataset(torchDataset):
     """
-    Dataset class to load a bunch of anndata from a lamin dataset (Collection) in a memory efficient way.
+    PyTorch Dataset for loading single-cell data from a LaminDB Collection.
 
-    This serves as a wrapper around lamin's mappedCollection to provide more features,
-    mostly, the management of hierarchical labels, the encoding of labels, the management of multiple species
+    This class wraps LaminDB's MappedCollection to provide additional features:
+    - Management of hierarchical ontology labels (cell type, tissue, disease, etc.)
+    - Automatic encoding of categorical labels to integers
+    - Multi-species gene handling with unified gene indexing
+    - Optional metacell aggregation and KNN neighbor retrieval
 
-    For an example of mappedDataset, see :meth:`~lamindb.Dataset.mapped`.
-
-    .. note::
-
-        A related data loader exists `here
-        <https://github.com/Genentech/scimilarity>`__.
+    The dataset lazily loads data from storage, making it memory-efficient for
+    large collections spanning multiple files.
 
     Args:
-    ----
-        lamin_dataset (lamindb.Dataset): lamin dataset to load
-        genedf (pd.Dataframe): dataframe containing the genes to load
-        obs (List[str]): list of observations to load from the Collection
-        clss_to_predict (List[str]): list of observations to encode
-        join_vars (flag): join variables @see :meth:`~lamindb.Dataset.mapped`.
-        hierarchical_clss: list of observations to map to a hierarchy using lamin's bionty
-        metacell_mode (float, optional): The mode to use for metacell sampling. Defaults to 0.0.
-        get_knn_cells (bool, optional): Whether to get the k-nearest neighbors of each cell. Defaults to False.
-        store_location (str, optional): The location to store the sampler indices. Defaults to None.
-        force_recompute_indices (bool, optional): Whether to force recompute the sampler indices. Defaults to False.
+        lamin_dataset (ln.Collection): LaminDB Collection containing the artifacts to load.
+        genedf (pd.DataFrame, optional): DataFrame with gene information, indexed by gene ID
+            with an 'organism' column. If None, automatically loaded based on organisms
+            in the dataset. Defaults to None.
+        clss_to_predict (List[str], optional): Observation columns to encode as prediction
+            targets. These will be integer-encoded in the output. Defaults to [].
+        hierarchical_clss (List[str], optional): Observation columns with hierarchical
+            ontology structure. These will have their ancestry relationships computed
+            using Bionty. Supported columns:
+            - "cell_type_ontology_term_id"
+            - "tissue_ontology_term_id"
+            - "disease_ontology_term_id"
+            - "development_stage_ontology_term_id"
+            - "assay_ontology_term_id"
+            - "self_reported_ethnicity_ontology_term_id"
+            Defaults to [].
+        join_vars (str, optional): How to join variables across artifacts.
+            "inner" for intersection, "outer" for union, None for no joining.
+            Defaults to None.
+        metacell_mode (float, optional): Probability of returning aggregated metacell
+            expression instead of single-cell. Defaults to 0.0.
+        get_knn_cells (bool, optional): Whether to include k-nearest neighbor cell
+            expression in the output. Requires precomputed neighbors in the data.
+            Defaults to False.
+        store_location (str, optional): Directory path to cache computed indices.
+            Defaults to None.
+        force_recompute_indices (bool, optional): Force recomputation of cached data.
+            Defaults to False.
+
+    Attributes:
+        mapped_dataset (MappedCollection): Underlying mapped collection for data access.
+        genedf (pd.DataFrame): Gene information DataFrame.
+        organisms (List[str]): List of organism ontology term IDs in the dataset.
+        class_topred (dict[str, set]): Mapping from class name to set of valid labels.
+        labels_groupings (dict[str, dict]): Hierarchical groupings for ontology classes.
+        encoder (dict[str, dict]): Label encoders mapping strings to integers.
+
+    Raises:
+        ValueError: If genedf is None and "organism_ontology_term_id" is not in clss_to_predict.
+
+    Example:
+        >>> collection = ln.Collection.filter(key="my_collection").first()
+        >>> dataset = Dataset(
+        ...     lamin_dataset=collection,
+        ...     clss_to_predict=["organism_ontology_term_id", "cell_type_ontology_term_id"],
+        ...     hierarchical_clss=["cell_type_ontology_term_id"],
+        ... )
+        >>> sample = dataset[0]  # Returns dict with "X" and encoded labels
     """
 
     lamin_dataset: ln.Collection
@@ -165,7 +201,20 @@ class Dataset(torchDataset):
         self,
         obs_keys: Union[str, List[str]],
     ):
-        """Get all categories for the given label keys."""
+        """
+        Get combined categorical codes for one or more label columns.
+
+        Retrieves labels from the mapped dataset and combines them into a single
+        categorical encoding. Useful for creating compound class labels for
+        stratified sampling.
+
+        Args:
+            obs_keys (str | List[str]): Column name(s) to retrieve and combine.
+
+        Returns:
+            np.ndarray: Integer codes representing the (combined) categories.
+                Shape: (n_samples,).
+        """
         if isinstance(obs_keys, str):
             obs_keys = [obs_keys]
         labels = None
@@ -179,25 +228,37 @@ class Dataset(torchDataset):
 
     def get_unseen_mapped_dataset_elements(self, idx: int):
         """
-        get_unseen_mapped_dataset_elements is a wrapper around mappedDataset.get_unseen_mapped_dataset_elements
+        Get genes marked as unseen for a specific sample.
+
+        Retrieves the list of genes that were not observed (expression = 0 or
+        marked as unseen) for the sample at the given index.
 
         Args:
-            idx (int): index of the element to get
+            idx (int): Sample index in the dataset.
 
         Returns:
-            List[str]: list of unseen genes
+            List[str]: List of unseen gene identifiers.
         """
         return [str(i)[2:-1] for i in self.mapped_dataset.uns(idx, "unseen_genes")]
 
     def define_hierarchies(self, clsses: List[str]):
         """
-        define_hierarchies is a method to define the hierarchies for the classes to predict
+        Define hierarchical label groupings from ontology relationships.
+
+        Uses Bionty to retrieve parent-child relationships for ontology terms,
+        then builds groupings mapping parent terms to their descendants.
+        Updates encoders to include parent terms and reorders labels so that
+        leaf terms (directly predictable) come first.
 
         Args:
-            clsses (List[str]): list of classes to predict
+            clsses (List[str]): List of ontology column names to process.
 
         Raises:
-            ValueError: if the class is not in the accepted classes
+            ValueError: If a class name is not in the supported ontology types.
+
+        Note:
+            Modifies self.labels_groupings, self.class_topred, and
+            self.mapped_dataset.encoders in place.
         """
         # TODO: use all possible hierarchies instead of just the ones for which we have a sample annotated with
         self.labels_groupings = {}
@@ -318,6 +379,7 @@ class Dataset(torchDataset):
 
 
 class SimpleAnnDataset(torchDataset):
+
     def __init__(
         self,
         adata: AnnData,
@@ -327,16 +389,40 @@ class SimpleAnnDataset(torchDataset):
         encoder: Optional[dict[str, dict]] = None,
     ):
         """
-        SimpleAnnDataset is a simple dataloader for an AnnData dataset. this is to interface nicely with the rest of
-        scDataloader and with your model during inference.
+        Simple PyTorch Dataset wrapper for a single AnnData object.
+
+        Provides a lightweight interface for using AnnData with PyTorch DataLoaders,
+        compatible with the scDataLoader collator. Useful for inference on new data
+        that isn't stored in LaminDB.
 
         Args:
-        ----
-            adata (anndata.AnnData): anndata object to use
-            obs_to_output (List[str]): list of observations to output from anndata.obs
-            layer (str): layer of the anndata to use
-            get_knn_cells (bool): whether to get the knn cells
-            encoder (dict[str, dict]): dictionary of encoders for the observations.
+            adata (AnnData): AnnData object containing expression data.
+            obs_to_output (List[str], optional): Observation columns to include in
+                output dictionaries. Defaults to [].
+            layer (str, optional): Layer name to use for expression values. If None,
+                uses adata.X. Defaults to None.
+            get_knn_cells (bool, optional): Whether to include k-nearest neighbor
+                expression data. Requires precomputed neighbors in adata.obsp.
+                Defaults to False.
+            encoder (dict[str, dict], optional): Dictionary mapping observation column
+                names to encoding dictionaries (str -> int). Defaults to None.
+
+        Attributes:
+            adataX (np.ndarray): Dense expression matrix.
+            encoder (dict): Label encoders.
+            obs_to_output (pd.DataFrame): Observation metadata to include.
+            distances (scipy.sparse matrix): KNN distance matrix (if get_knn_cells=True).
+
+        Raises:
+            ValueError: If get_knn_cells=True but "connectivities" not in adata.obsp.
+
+        Example:
+            >>> dataset = SimpleAnnDataset(
+            ...     adata=my_adata,
+            ...     obs_to_output=["cell_type", "organism_ontology_term_id"],
+            ...     encoder={"organism_ontology_term_id": {"NCBITaxon:9606": 0}},
+            ... )
+            >>> loader = DataLoader(dataset, batch_size=32, collate_fn=collator)
         """
         self.adataX = adata.layers[layer] if layer is not None else adata.X
         self.adataX = self.adataX.toarray() if issparse(self.adataX) else self.adataX
@@ -392,6 +478,46 @@ def mapped(
     store_location: str | None = None,
     force_recompute_indices: bool = False,
 ) -> MappedCollection:
+    """
+    Create a MappedCollection from a LaminDB Collection.
+
+    Factory function that handles artifact path resolution (staging or streaming)
+    and creates a MappedCollection for efficient access to multiple h5ad/zarr files.
+
+    Args:
+        dataset (ln.Collection): LaminDB Collection containing artifacts to map.
+        obs_keys (List[str], optional): Observation columns to load. Defaults to None.
+        obsm_keys (List[str], optional): Obsm keys to load. Defaults to None.
+        obs_filter (dict, optional): Filter observations by column values.
+            Keys are column names, values are allowed values. Defaults to None.
+        join (str, optional): How to join variables across files. "inner" for
+            intersection, "outer" for union. Defaults to "inner".
+        encode_labels (bool | List[str], optional): Whether/which columns to
+            integer-encode. True encodes all obs_keys. Defaults to True.
+        unknown_label (str | dict, optional): Label to use for unknown/missing
+            categories. Defaults to None.
+        cache_categories (bool, optional): Whether to cache category mappings.
+            Defaults to True.
+        parallel (bool, optional): Enable parallel data loading. Defaults to False.
+        dtype (str, optional): Data type for expression values. Defaults to None.
+        stream (bool, optional): If True, stream from cloud storage instead of
+            staging locally. Defaults to False.
+        is_run_input (bool, optional): Track as run input in LaminDB. Defaults to None.
+        metacell_mode (bool, optional): Enable metacell aggregation. Defaults to False.
+        meta_assays (List[str], optional): Assay types to treat as metacell-like.
+            Defaults to ["EFO:0022857", "EFO:0010961"].
+        get_knn_cells (bool, optional): Include KNN neighbor data. Defaults to False.
+        store_location (str, optional): Cache directory path. Defaults to None.
+        force_recompute_indices (bool, optional): Force recompute cached data.
+            Defaults to False.
+
+    Returns:
+        MappedCollection: Mapped collection for data access.
+
+    Note:
+        Artifacts with suffixes other than .h5ad, .zrad, .zarr are ignored.
+        Non-existent paths are skipped with a warning.
+    """
     path_list = []
     for artifact in dataset.artifacts.all():
         if artifact.suffix not in {".h5ad", ".zrad", ".zarr"}:
@@ -425,14 +551,26 @@ def mapped(
 
 
 def concat_categorical_codes(series_list: List[pd.Categorical]) -> pd.Categorical:
-    """Efficiently combine multiple categorical data using their codes,
-    only creating categories for combinations that exist in the data.
+    """
+    Efficiently combine multiple categorical arrays into a single encoding.
+
+    Creates a combined categorical where each unique combination of input
+    categories gets a unique code. Only combinations that exist in the data
+    are assigned codes (sparse encoding).
 
     Args:
-        series_list: List of pandas Categorical data
+        series_list (List[pd.Categorical]): List of categorical arrays to combine.
+            All arrays must have the same length.
 
     Returns:
-        Combined Categorical with only existing combinations
+        pd.Categorical: Combined categorical with compressed codes representing
+            unique combinations present in the data.
+
+    Example:
+        >>> cat1 = pd.Categorical(["a", "a", "b", "b"])
+        >>> cat2 = pd.Categorical(["x", "y", "x", "y"])
+        >>> combined = concat_categorical_codes([cat1, cat2])
+        >>> # Results in 4 unique codes for (a,x), (a,y), (b,x), (b,y)
     """
     # Get the codes for each categorical
     codes_list = [s.codes.astype(np.int32) for s in series_list]

{scdataloader-2.0.2 → scdataloader-2.0.4}/scdataloader/datamodule.py

@@ -65,42 +65,102 @@ class DataModule(L.LightningDataModule):
         genedf: Optional[pd.DataFrame] = None,
         n_bins: int = 0,
         curiculum: int = 0,
+        start_at: int = 0,
         **kwargs,
     ):
         """
-        DataModule a pytorch lighting datamodule directly from a lamin Collection.
-        it can work with bare pytorch too
+        PyTorch Lightning DataModule for loading single-cell data from a LaminDB Collection.
 
-        It implements train / val / test dataloaders. the train is weighted random, val is random, test is one to many separated datasets.
-        This is where the mappedCollection, dataset, and collator are combined to create the dataloaders.
+        This DataModule provides train/val/test dataloaders with configurable sampling strategies.
+        It combines MappedCollection, Dataset, and Collator to create efficient data pipelines
+        for training single-cell foundation models.
+
+        The training dataloader uses weighted random sampling based on class frequencies,
+        validation uses random sampling, and test uses sequential sampling on held-out datasets.
 
         Args:
-            collection_name (str): The lamindb collection to be used.
-            weight_scaler (int, optional): how much more you will see the most present vs less present category.
-            n_samples_per_epoch (int, optional): The number of samples to include in the training set for each epoch. Defaults to 2_000_000.
-            validation_split (float, optional): The proportion of the dataset to include in the validation split. Defaults to 0.2.
-            test_split (float, optional): The proportion of the dataset to include in the test split. Defaults to 0.
-                it will use a full dataset and will round to the nearest dataset's cell count.
-            use_default_col (bool, optional): Whether to use the default collator. Defaults to True.
-            clss_to_weight (List[str], optional): List of labels to weight in the trainer's weighted random sampler. Defaults to [].
-            assays_to_drop (List[str], optional): List of assays to drop from the dataset. Defaults to [].
-            gene_pos_file (Union[bool, str], optional): The path to the gene positions file. Defaults to True.
-                the file must have ensembl_gene_id as index.
-                This is used to subset the available genes further to the ones that have embeddings in your model.
-            max_len (int, optional): The maximum length of the input tensor. Defaults to 1000.
-            how (str, optional): The method to use for the collator. Defaults to "random expr".
-            organism_col (str, optional): The name of the organism. Defaults to "organism_ontology_term_id".
-            tp_name (Optional[str], optional): The name of the timepoint. Defaults to None.
-            hierarchical_clss (List[str], optional): List of hierarchical classes. Defaults to [].
-            metacell_mode (float, optional): The probability of using metacell mode. Defaults to 0.0.
-            clss_to_predict (List[str], optional): List of classes to predict. Defaults to ["organism_ontology_term_id"].
-            get_knn_cells (bool, optional): Whether to get the k-nearest neighbors of each queried cells. Defaults to False.
-            store_location (str, optional): The location to store the sampler indices. Defaults to None.
-            force_recompute_indices (bool, optional): Whether to force recompute the sampler indices. Defaults to False.
-            sampler_workers (int, optional): The number of workers to use for the sampler. Defaults to None (auto-determined).
-            sampler_chunk_size (int, optional): The size of the chunks to use for the sampler. Defaults to None (auto-determined).
-            **kwargs: Additional keyword arguments passed to the pytorch DataLoader.
-            see @file data.py and @file collator.py for more details about some of the parameters
+            collection_name (str): Key of the LaminDB Collection to load.
+            clss_to_weight (List[str], optional): Label columns to use for weighted sampling
+                in the training dataloader. Supports "nnz" for weighting by number of
+                non-zero genes. Defaults to ["organism_ontology_term_id"].
+            weight_scaler (int, optional): Controls balance between rare and common classes.
+                Higher values lead to more uniform sampling across classes. Set to 0 to
+                disable weighted sampling. Defaults to 10.
+            n_samples_per_epoch (int, optional): Number of samples to draw per training epoch.
+                Defaults to 2,000,000.
+            validation_split (float | int, optional): Proportion (float) or absolute number (int)
+                of samples for validation. Defaults to 0.2.
+            test_split (float | int, optional): Proportion (float) or absolute number (int)
+                of samples for testing. Uses entire datasets as test sets, rounding to
+                nearest dataset boundary. Defaults to 0.
+            use_default_col (bool, optional): Whether to use the default Collator for batch
+                preparation. If False, no collate_fn is applied. Defaults to True.
+            clss_to_predict (List[str], optional): Observation columns to encode as prediction
+                targets. Must include "organism_ontology_term_id". Defaults to
+                ["organism_ontology_term_id"].
+            hierarchical_clss (List[str], optional): Observation columns with hierarchical
+                ontology structure to be processed. Defaults to [].
+            how (str, optional): Gene selection strategy passed to Collator. One of
+                "most expr", "random expr", "all", "some". Defaults to "random expr".
+            organism_col (str, optional): Column name for organism ontology term ID.
+                Defaults to "organism_ontology_term_id".
+            max_len (int, optional): Maximum number of genes per sample passed to Collator.
+                Defaults to 1000.
+            replacement (bool, optional): Whether to sample with replacement in training.
+                Defaults to True.
+            gene_subset (List[str], optional): List of genes to restrict the dataset to.
+                Useful when model only supports specific genes. Defaults to None.
+            tp_name (str, optional): Column name for time point or heat diffusion values.
+                Defaults to None.
+            assays_to_drop (List[str], optional): List of assay ontology term IDs to exclude
+                from training. Defaults to ["EFO:0030007"] (ATAC-seq).
+            metacell_mode (float, optional): Probability of using metacell aggregation mode.
+                Cannot be used with get_knn_cells. Defaults to 0.0.
+            get_knn_cells (bool, optional): Whether to include k-nearest neighbor cell
+                expression data. Cannot be used with metacell_mode. Defaults to False.
+            store_location (str, optional): Directory path to cache sampler indices and
+                labels for faster subsequent loading. Defaults to None.
+            force_recompute_indices (bool, optional): Force recomputation of cached indices
+                even if they exist. Defaults to False.
+            sampler_workers (int, optional): Number of parallel workers for building sampler
+                indices. Auto-determined based on available CPUs if None. Defaults to None.
+            sampler_chunk_size (int, optional): Chunk size for parallel sampler processing.
+                Auto-determined based on available memory if None. Defaults to None.
+            organisms (List[str], optional): List of organisms to include. If None, uses
+                all organisms in the dataset. Defaults to None.
+            genedf (pd.DataFrame, optional): Gene information DataFrame. If None, loaded
+                automatically. Defaults to None.
+            n_bins (int, optional): Number of bins for expression discretization. 0 means
+                no binning. Defaults to 0.
+            curiculum (int, optional): Curriculum learning parameter. If > 0, gradually
+                increases sampling weight balance over epochs. Defaults to 0.
+            start_at (int, optional): Starting index for resuming inference. Requires same
+                number of GPUs as previous run. Defaults to 0.
+            **kwargs: Additional arguments passed to PyTorch DataLoader (e.g., batch_size,
+                num_workers, pin_memory).
+
+        Attributes:
+            dataset (Dataset): The underlying Dataset instance.
+            classes (dict[str, int]): Mapping from class names to number of categories.
+            train_labels (np.ndarray): Label array for weighted sampling.
+            idx_full (np.ndarray): Indices for training samples.
+            valid_idx (np.ndarray): Indices for validation samples.
+            test_idx (np.ndarray): Indices for test samples.
+            test_datasets (List[str]): Paths to datasets used for testing.
+
+        Raises:
+            ValueError: If "organism_ontology_term_id" not in clss_to_predict.
+            ValueError: If both metacell_mode > 0 and get_knn_cells are True.
+
+        Example:
+            >>> dm = DataModule(
+            ...     collection_name="my_collection",
+            ...     batch_size=32,
+            ...     num_workers=4,
+            ...     max_len=2000,
+            ... )
+            >>> dm.setup()
+            >>> train_loader = dm.train_dataloader()
         """
         if "organism_ontology_term_id" not in clss_to_predict:
             raise ValueError(
@@ -162,6 +222,7 @@ class DataModule(L.LightningDataModule):
         self.sampler_chunk_size = sampler_chunk_size
         self.store_location = store_location
         self.nnz = None
+        self.start_at = start_at
         self.idx_full = None
         self.max_len = max_len
         self.test_datasets = []
@@ -283,12 +344,24 @@ class DataModule(L.LightningDataModule):
 
     def setup(self, stage=None):
         """
-        setup method is used to prepare the data for the training, validation, and test sets.
-        It shuffles the data, calculates weights for each set, and creates samplers for each set.
+        Prepare data splits for training, validation, and testing.
+
+        This method shuffles the data, computes sample weights for weighted sampling,
+        removes samples from dropped assays, and creates train/val/test splits.
+        Test splits use entire datasets to ensure evaluation on unseen data sources.
+
+        Results can be cached to `store_location` for faster subsequent runs.
 
         Args:
-            stage (str, optional): The stage of the model training process.
-            It can be either 'fit' or 'test'. Defaults to None.
+            stage (str, optional): Training stage ('fit', 'test', or None for both).
+                Currently not used but kept for Lightning compatibility. Defaults to None.
+
+        Returns:
+            List[str]: List of paths to test datasets.
+
+        Note:
+            Must be called before using dataloaders. The train/val/test split is
+            deterministic when loading from cache.
         """
         print("setting up the datamodule")
         start_time = time.time()
@@ -324,9 +397,9 @@ class DataModule(L.LightningDataModule):
                 len_test = self.test_split
             else:
                 len_test = int(self.n_samples * self.test_split)
-            assert len_test + len_valid < self.n_samples, (
-                "test set + valid set size is configured to be larger than entire dataset."
-            )
+            assert (
+                len_test + len_valid < self.n_samples
+            ), "test set + valid set size is configured to be larger than entire dataset."
 
             idx_full = []
             if len(self.assays_to_drop) > 0:
@@ -439,6 +512,22 @@ class DataModule(L.LightningDataModule):
         return self.test_datasets
 
     def train_dataloader(self, **kwargs):
+        """
+        Create the training DataLoader with weighted random sampling.
+
+        Uses LabelWeightedSampler for class-balanced sampling when weight_scaler > 0
+        and clss_to_weight is specified. Otherwise uses RankShardSampler for
+        distributed training without weighting.
+
+        Args:
+            **kwargs: Additional arguments passed to DataLoader, overriding defaults.
+
+        Returns:
+            DataLoader: Training DataLoader instance.
+
+        Raises:
+            ValueError: If setup() has not been called.
+        """
         if len(self.clss_to_weight) > 0 and self.weight_scaler > 0:
             try:
                 print("Setting up the parallel train sampler...")
@@ -461,7 +550,7 @@ class DataModule(L.LightningDataModule):
             dataset = None
         else:
             dataset = Subset(self.dataset, self.idx_full)
-            train_sampler = RankShardSampler(len(dataset))
+            train_sampler = RankShardSampler(len(dataset), start_at=self.start_at)
         current_loader_kwargs = kwargs.copy()
         current_loader_kwargs.update(self.kwargs)
         return DataLoader(
@@ -471,6 +560,12 @@ class DataModule(L.LightningDataModule):
         )
 
     def val_dataloader(self):
+        """
+        Create the validation DataLoader.
+
+        Returns:
+            DataLoader | List: Validation DataLoader, or empty list if no validation split.
+        """
         return (
             DataLoader(
                 Subset(self.dataset, self.valid_idx),
@@ -481,6 +576,12 @@ class DataModule(L.LightningDataModule):
         )
 
     def test_dataloader(self):
+        """
+        Create the test DataLoader with sequential sampling.
+
+        Returns:
+            DataLoader | List: Test DataLoader, or empty list if no test split.
+        """
         return (
             DataLoader(
                 self.dataset, sampler=SequentialSampler(self.test_idx), **self.kwargs
@@ -490,24 +591,23 @@ class DataModule(L.LightningDataModule):
         )
 
     def predict_dataloader(self):
+        """
+        Create a DataLoader for prediction over all training data.
+
+        Uses RankShardSampler for distributed inference.
+
+        Returns:
+            DataLoader: Prediction DataLoader instance.
+        """
         subset = Subset(self.dataset, self.idx_full)
         return DataLoader(
-            subset,
-            sampler=RankShardSampler(len(subset)),
+            self.dataset,
+            sampler=RankShardSampler(len(subset), start_at=self.start_at),
             **self.kwargs,
         )
 
 
 class LabelWeightedSampler(Sampler[int]):
-    """
-    A weighted random sampler that samples from a dataset with respect to both class weights and element weights.
-
-    This sampler is designed to handle very large datasets efficiently, with optimizations for:
-    1. Parallel building of class indices
-    2. Chunked processing for large arrays
-    3. Efficient memory management
-    4. Proper handling of replacement and non-replacement sampling
-    """
 
     label_weights: torch.Tensor
     klass_indices: dict[int, torch.Tensor]
@@ -529,16 +629,58 @@ class LabelWeightedSampler(Sampler[int]):
         curiculum: int = 0,
     ) -> None:
         """
-        Initialize the sampler with parallel processing for large datasets.
+        Weighted random sampler balancing both class frequencies and element weights.
+
+        This sampler is optimized for very large datasets (millions of samples) with:
+        - Parallel construction of class indices using multiple CPU workers
+        - Chunked processing to manage memory usage
+        - Support for curriculum learning via progressive weight scaling
+        - Optional per-element weights (e.g., by number of expressed genes)
+
+        The sampling process:
+        1. Sample class labels according to class weights
+        2. For each sampled class, sample elements according to element weights
+        3. Shuffle all sampled indices
 
         Args:
-            weight_scaler: Scaling factor for class weights (higher means less balanced sampling)
-            labels: Class label for each dataset element (length = dataset size)
-            num_samples: Number of samples to draw
-            replacement: Whether to sample with replacement
-            element_weights: Optional weights for each element within classes
-            n_workers: Number of parallel workers to use (default: number of CPUs-1)
-            chunk_size: Size of chunks to process in parallel (default: 10M elements)
+            labels (np.ndarray): Integer class label for each dataset element.
+                Shape: (dataset_size,). The last unique label is treated as
+                "excluded" with weight 0.
+            num_samples (int): Number of samples to draw per epoch.
+            replacement (bool, optional): Whether to sample with replacement.
+                Defaults to True.
+            weight_scaler (float, optional): Controls class weight balance.
+                Weight formula: (scaler * count) / (count + scaler).
+                Higher values = more uniform sampling. Defaults to None.
+            element_weights (Sequence[float], optional): Per-element sampling weights.
+                Shape: (dataset_size,). Defaults to None (uniform within class).
+            n_workers (int, optional): Number of parallel workers for index building.
+                Defaults to min(20, num_cpus - 1).
+            chunk_size (int, optional): Elements per chunk for parallel processing.
+                Auto-determined based on available memory if None.
+            store_location (str, optional): Directory to cache computed indices.
+                Defaults to None.
+            force_recompute_indices (bool, optional): Recompute indices even if cached.
+                Defaults to False.
+            curiculum (int, optional): Curriculum learning epochs. If > 0, weight
+                exponent increases from 0 to 1 over this many epochs. Defaults to 0.
+
+        Attributes:
+            label_weights (torch.Tensor): Computed weights per class label.
+            klass_indices (torch.Tensor): Concatenated indices for all classes.
+            klass_offsets (torch.Tensor): Starting offset for each class in klass_indices.
+            count (int): Number of times __iter__ has been called (for curriculum).
+
+        Example:
+            >>> sampler = LabelWeightedSampler(
+            ...     labels=train_labels,
+            ...     num_samples=1_000_000,
+            ...     weight_scaler=10,
+            ...     element_weights=nnz_weights,
+            ... )
+            >>> for idx in sampler:
+            ...     # Process sample at idx
+            ...     pass
         """
         print("Initializing optimized parallel weighted sampler...")
         super(LabelWeightedSampler, self).__init__(None)
@@ -667,7 +809,9 @@ class LabelWeightedSampler(Sampler[int]):
         unique_samples, sample_counts = torch.unique(sample_labels, return_counts=True)
 
         # Initialize result tensor
-        result_indices_list = []  # Changed name to avoid conflict if you had result_indices elsewhere
+        result_indices_list = (
+            []
+        )  # Changed name to avoid conflict if you had result_indices elsewhere
 
         # Process only the classes that were actually sampled
         for i, (label, count) in tqdm(
@@ -847,11 +991,36 @@ class LabelWeightedSampler(Sampler[int]):
 
 
 class RankShardSampler(Sampler[int]):
-    """Shards a dataset contiguously across ranks without padding or duplicates.
-    Preserves the existing order (e.g., your pre-shuffled idx_full)."""
+    """
+    Sampler that shards data contiguously across distributed ranks.
+
+    Divides the dataset into contiguous chunks, one per rank, without
+    padding or duplicating samples. Preserves the original data order
+    within each shard (useful for pre-shuffled data).
+
+    Args:
+        data_len (int): Total number of samples in the dataset.
+        start_at (int, optional): Global starting index for resuming training.
+            Requires the same number of GPUs as the previous run. Defaults to 0.
+
+    Attributes:
+        rank (int): Current process rank (0 if not distributed).
+        world_size (int): Total number of processes (1 if not distributed).
+        start (int): Starting index for this rank's shard.
+        end (int): Ending index (exclusive) for this rank's shard.
+
+    Note:
+        The last rank may have fewer samples than others if the dataset
+        size is not evenly divisible by world_size.
+
+    Example:
+        >>> sampler = RankShardSampler(len(dataset))
+        >>> loader = DataLoader(dataset, sampler=sampler)
+    """
 
-    def __init__(self, data_len: int):
+    def __init__(self, data_len: int, start_at: int = 0) -> None:
         self.data_len = data_len
+        self.start_at = start_at
         if torch.distributed.is_available() and torch.distributed.is_initialized():
             self.rank = torch.distributed.get_rank()
             self.world_size = torch.distributed.get_world_size()
@@ -859,9 +1028,16 @@ class RankShardSampler(Sampler[int]):
             self.rank, self.world_size = 0, 1
 
         # contiguous chunk per rank (last rank may be shorter)
+        if self.start_at > 0:
+            print(
+                "!!!!ATTENTION: make sure that you are running on the exact same \
+                number of GPU as your previous run!!!!!"
+            )
+        print(f"Sharding data of size {data_len} over {self.world_size} ranks")
         per_rank = math.ceil(self.data_len / self.world_size)
-        self.start = self.rank * per_rank
-        self.end = min(self.start + per_rank, self.data_len)
+        self.start = int((self.start_at / self.world_size) + (self.rank * per_rank))
+        self.end = min((self.rank + 1) * per_rank, self.data_len)
+        print(f"Rank {self.rank} processing indices from {self.start} to {self.end}")
 
     def __iter__(self):
         return iter(range(self.start, self.end))

{scdataloader-2.0.2 → scdataloader-2.0.4}/scdataloader/preprocess.py

@@ -1,3 +1,10 @@
+"""
+Preprocessing utilities for single-cell gene expression data.
+
+This module provides functions for normalizing, transforming, and discretizing
+gene expression values for use with scPRINT and similar models.
+"""
+
 import gc
 import time
 from typing import Callable, List, Optional, Union
@@ -664,39 +671,41 @@ def is_log1p(adata: AnnData) -> bool:
     return True
 
 
-def _digitize(x: np.ndarray, bins: np.ndarray, side="both") -> np.ndarray:
+def _digitize(values: np.ndarray, bins: np.ndarray) -> np.ndarray:
     """
-    Digitize the data into bins. This method spreads data uniformly when bins
-    have same values.
+    Digitize values into discrete bins with 1-based indexing.
 
-    Args:
+    Similar to np.digitize but ensures output is 1-indexed (bin 0 reserved for
+    zero values) and handles edge cases for expression binning.
 
-    x (:class:`np.ndarray`):
-        The data to digitize.
-    bins (:class:`np.ndarray`):
-        The bins to use for digitization, in increasing order.
-    side (:class:`str`, optional):
-        The side to use for digitization. If "one", the left side is used. If
-        "both", the left and right side are used. Default to "one".
+    Args:
+        values (np.ndarray): Array of values to discretize. Should be non-zero
+            expression values.
+        bins (np.ndarray): Bin edges from np.quantile or similar. Values are
+            assigned to bins based on which edges they fall between.
 
     Returns:
-
-    :class:`np.ndarray`:
-        The digitized data.
+        np.ndarray: Integer bin indices, 1-indexed. Values equal to bins[i]
+            are assigned to bin i+1.
+
+    Example:
+        >>> values = np.array([0.5, 1.5, 2.5, 3.5])
+        >>> bins = np.array([1.0, 2.0, 3.0])
+        >>> _digitize(values, bins)
+        array([1, 2, 3, 3])
+
+    Note:
+        This function is used internally by the Collator for expression binning.
+        Zero values should be handled separately before calling this function.
     """
-    assert x.ndim == 1 and bins.ndim == 1
-
-    left_digits = np.digitize(x, bins)
-    if side == "one":
-        return left_digits
+    assert values.ndim == 1 and bins.ndim == 1
 
-    right_difits = np.digitize(x, bins, right=True)
+    left_digits = np.digitize(values, bins)
+    return left_digits
 
-    rands = np.random.rand(len(x))  # uniform random numbers
 
-    digits = rands * (right_difits - left_digits) + left_digits
-    digits = np.ceil(digits).astype(np.int64)
-    return digits
+# Add documentation for any other functions in preprocess.py
+# ...existing code...
 
 
 def binning(row: np.ndarray, n_bins: int) -> np.ndarray: