scdataloader 1.8.0__py3-none-any.whl → 1.8.1__py3-none-any.whl

scdataloader/VERSION

@@ -1 +1 @@
-1.8.0
+1.8.1

scdataloader/__init__.py

@@ -2,5 +2,6 @@ from .collator import Collator
 from .data import Dataset, SimpleAnnDataset
 from .datamodule import DataModule
 from .preprocess import Preprocessor
+from importlib.metadata import version
 
-__version__ = "1.7.0"
+__version__ = version("scdataloader")

scdataloader/collator.py

@@ -24,7 +24,6 @@ class Collator:
         genelist: list[str] = [],
         downsample: Optional[float] = None,  # don't use it for training!
         save_output: Optional[str] = None,
-        metacell_mode: bool = False,
     ):
         """
         This class is responsible for collating data for the scPRINT model. It handles the
@@ -62,7 +61,6 @@ class Collator:
                 This is usually done by the scPRINT model during training but this option allows you to do it directly from the collator
             save_output (str, optional): If not None, saves the output to a file. Defaults to None.
                 This is mainly for debugging purposes
-            metacell_mode (bool, optional): Whether to sample a metacell. Defaults to False.
         """
         self.organisms = organisms
         self.genedf = load_genes(organisms)
@@ -82,7 +80,6 @@ class Collator:
         self.accepted_genes = {}
         self.downsample = downsample
         self.to_subset = {}
-        self.metacell_mode = metacell_mode
         self._setup(org_to_id, valid_genes, genelist)
 
     def _setup(self, org_to_id=None, valid_genes=[], genelist=[]):
@@ -135,6 +132,7 @@ class Collator:
         dataset = []
         nnz_loc = []
         is_meta = []
+        knn_cells = []
         for elem in batch:
             organism_id = elem[self.organism_name]
             if organism_id not in self.organism_ids:
@@ -145,14 +143,24 @@ class Collator:
             total_count.append(expr.sum())
             if len(self.accepted_genes) > 0:
                 expr = expr[self.accepted_genes[organism_id]]
+                if "knn_cells" in elem:
+                    elem["knn_cells"] = elem["knn_cells"][
+                        :, self.accepted_genes[organism_id]
+                    ]
             if self.how == "most expr":
-                nnz_loc = np.where(expr > 0)[0]
+                if "knn_cells" in elem:
+                    nnz_loc = np.where(expr + elem["knn_cells"].sum(0) > 0)[0]
+                else:
+                    nnz_loc = np.where(expr > 0)[0]
                 ma = self.max_len if self.max_len < len(nnz_loc) else len(nnz_loc)
                 loc = np.argsort(expr)[-(ma):][::-1]
                 # nnz_loc = [1] * 30_000
                 # loc = np.argsort(expr)[-(self.max_len) :][::-1]
             elif self.how == "random expr":
-                nnz_loc = np.where(expr > 0)[0]
+                if "knn_cells" in elem:
+                    nnz_loc = np.where(expr + elem["knn_cells"].sum(0) > 0)[0]
+                else:
+                    nnz_loc = np.where(expr > 0)[0]
                 loc = nnz_loc[
                     np.random.choice(
                         len(nnz_loc),
@@ -171,7 +179,10 @@ class Collator:
                 "all",
                 "some",
             ]:
-                zero_loc = np.where(expr == 0)[0]
+                if "knn_cells" in elem:
+                    zero_loc = np.where(expr + elem["knn_cells"].sum(0) == 0)[0]
+                else:
+                    zero_loc = np.where(expr == 0)[0]
                 zero_loc = zero_loc[
                     np.random.choice(
                         len(zero_loc),
@@ -185,9 +196,13 @@ class Collator:
                     )
                 ]
                 loc = np.concatenate((loc, zero_loc), axis=None)
+            if "knn_cells" in elem:
+                knn_cells.append(elem["knn_cells"][:, loc])
             expr = expr[loc]
             loc = loc + self.start_idx[organism_id]
             if self.how == "some":
+                if "knn_cells" in elem:
+                    knn_cells[-1] = knn_cells[-1][self.to_subset[organism_id]]
                 expr = expr[self.to_subset[organism_id]]
                 loc = loc[self.to_subset[organism_id]]
             exprs.append(expr)
@@ -197,7 +212,7 @@ class Collator:
                 tp.append(elem[self.tp_name])
             else:
                 tp.append(0)
-            if self.metacell_mode:
+            if "is_meta" in elem:
                 is_meta.append(elem["is_meta"])
             other_classes.append([elem[i] for i in self.class_names])
         expr = np.array(exprs)
@@ -207,6 +222,7 @@ class Collator:
         other_classes = np.array(other_classes)
         dataset = np.array(dataset)
         is_meta = np.array(is_meta)
+        knn_cells = np.array(knn_cells)
         # normalize counts
         if self.norm_to is not None:
             expr = (expr * self.norm_to) / total_count[:, None]
@@ -217,15 +233,6 @@ class Collator:
         if self.n_bins:
             pass
 
-        # find the associated gene ids (given the species)
-
-        # get the NN cells
-
-        # do encoding / selection a la scGPT
-
-        # do encoding of graph location
-        # encode all the edges in some sparse way
-        # normalizing total counts between 0,1
         ret = {
             "x": Tensor(expr),
             "genes": Tensor(gene_locs).int(),
@@ -233,8 +240,10 @@ class Collator:
             "tp": Tensor(tp),
             "depth": Tensor(total_count),
         }
-        if self.metacell_mode:
+        if len(is_meta) > 0:
             ret.update({"is_meta": Tensor(is_meta).int()})
+        if len(knn_cells) > 0:
+            ret.update({"knn_cells": Tensor(knn_cells).int()})
         if len(dataset) > 0:
             ret.update({"dataset": Tensor(dataset).to(long)})
         if self.downsample is not None:

scdataloader/data.py

@@ -58,6 +58,7 @@ class Dataset(torchDataset):
     hierarchical_clss: Optional[list[str]] = field(default_factory=list)
     join_vars: Literal["inner", "outer"] | None = None
     metacell_mode: float = 0.0
+    get_knn_cells: bool = False
 
     def __post_init__(self):
         self.mapped_dataset = mapped(
@@ -69,6 +70,7 @@ class Dataset(torchDataset):
             stream=True,
             parallel=True,
             metacell_mode=self.metacell_mode,
+            get_knn_cells=self.get_knn_cells,
         )
         print(
             "won't do any check but we recommend to have your dataset coming from local storage"
@@ -371,6 +373,7 @@ def mapped(
     is_run_input: bool | None = None,
     metacell_mode: bool = False,
     meta_assays: list[str] = ["EFO:0022857", "EFO:0010961"],
+    get_knn_cells: bool = False,
 ) -> MappedCollection:
     path_list = []
     for artifact in dataset.artifacts.all():
@@ -397,5 +400,6 @@ def mapped(
         dtype=dtype,
         meta_assays=meta_assays,
         metacell_mode=metacell_mode,
+        get_knn_cells=get_knn_cells,
     )
     return ds

scdataloader/datamodule.py

@@ -52,6 +52,7 @@ class DataModule(L.LightningDataModule):
             # "EFO:0030062", # slide-seq
         ],
         metacell_mode: float = 0.0,
+        get_knn_cells: bool = False,
         modify_seed_on_requeue: bool = True,
         **kwargs,
     ):
@@ -88,6 +89,7 @@ class DataModule(L.LightningDataModule):
             metacell_mode (float, optional): The probability of using metacell mode. Defaults to 0.0.
             clss_to_predict (list, optional): List of classes to predict. Defaults to ["organism_ontology_term_id"].
             modify_seed_on_requeue (bool, optional): Whether to modify the seed on requeue. Defaults to True.
+            get_knn_cells (bool, optional): Whether to get the k-nearest neighbors of each queried cells. Defaults to False.
             **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
         """
@@ -98,6 +100,7 @@ class DataModule(L.LightningDataModule):
                 clss_to_predict=clss_to_predict,
                 hierarchical_clss=hierarchical_clss,
                 metacell_mode=metacell_mode,
+                get_knn_cells=get_knn_cells,
             )
         # and location
         self.metacell_mode = bool(metacell_mode)
@@ -157,7 +160,6 @@ class DataModule(L.LightningDataModule):
                 tp_name=tp_name,
                 organism_name=organism_name,
                 class_names=clss_to_predict,
-                metacell_mode=bool(metacell_mode),
             )
         self.validation_split = validation_split
         self.test_split = test_split

scdataloader/mapped.py

@@ -96,8 +96,9 @@ class MappedCollection:
         cache_categories: Enable caching categories of ``obs_keys`` for faster access.
         parallel: Enable sampling with multiple processes.
         dtype: Convert numpy arrays from ``.X``, ``.layers`` and ``.obsm``
-        meta_assays: Assays to check for metacells.
-        metacell_mode: Mode for metacells.
+        meta_assays: Assays that are already defined as metacells.
+        metacell_mode: frequency at which to sample a metacell (an average of k-nearest neighbors).
+        get_knn_cells: Whether to also dataload the k-nearest neighbors of each queried cells.
     """
 
     def __init__(
@@ -114,6 +115,7 @@ class MappedCollection:
         parallel: bool = False,
         dtype: str | None = None,
         metacell_mode: float = 0.0,
+        get_knn_cells: bool = False,
         meta_assays: list[str] = ["EFO:0022857", "EFO:0010961"],
     ):
         if join not in {None, "inner", "outer"}:  # pragma: nocover
@@ -166,6 +168,7 @@ class MappedCollection:
         self.metacell_mode = metacell_mode
         self.path_list = path_list
         self.meta_assays = meta_assays
+        self.get_knn_cells = get_knn_cells
         self._make_connections(path_list, parallel)
 
         self._cache_cats: dict = {}
@@ -396,12 +399,15 @@ class MappedCollection:
                         label_idx = self.encoders[label][label_idx]
                     out[label] = label_idx
 
-            out["is_meta"] = False
-            if len(self.meta_assays) > 0 and "assay_ontology_term_id" in self.obs_keys:
-                if out["assay_ontology_term_id"] in self.meta_assays:
-                    out["is_meta"] = True
-                    return out
             if self.metacell_mode > 0:
+                if (
+                    len(self.meta_assays) > 0
+                    and "assay_ontology_term_id" in self.obs_keys
+                ):
+                    if out["assay_ontology_term_id"] in self.meta_assays:
+                        out["is_meta"] = True
+                        return out
+                out["is_meta"] = False
                 if np.random.random() < self.metacell_mode:
                     out["is_meta"] = True
                     distances = self._get_data_idx(store["obsp"]["distances"], obs_idx)
@@ -410,6 +416,18 @@ class MappedCollection:
                         out[layers_key] += self._get_data_idx(
                             lazy_data, i, self.join_vars, var_idxs_join, self.n_vars
                         )
+            elif self.get_knn_cells:
+                distances = self._get_data_idx(store["obsp"]["distances"], obs_idx)
+                nn_idx = np.argsort(-1 / (distances - 1e-6))[:6]
+                out["knn_cells"] = np.array(
+                    [
+                        self._get_data_idx(
+                            lazy_data, i, self.join_vars, var_idxs_join, self.n_vars
+                        )
+                        for i in nn_idx
+                    ],
+                    dtype=int,
+                )
 
         return out
 

scdataloader/preprocess.py

@@ -64,6 +64,11 @@ class Preprocessor:
         """
         Initializes the preprocessor and configures the workflow steps.
 
+        Your dataset should contain at least the following obs:
+        - `organism_ontology_term_id` with the ontology id of the organism of your anndata
+        - gene names in the `var.index` field of your anndata that map to the ensembl_gene nomenclature
+        or the hugo gene symbols nomenclature (if the later, set `is_symbol` to True)
+
         Args:
             filter_gene_by_counts (int or bool, optional): Determines whether to filter genes by counts.
                 If int, filters genes with counts. Defaults to False.
@@ -130,6 +135,14 @@ class Preprocessor:
         self.keepdata = keepdata
 
     def __call__(self, adata, dataset_id=None) -> AnnData:
+        if "organism_ontology_term_id" not in adata[0].obs.columns:
+            raise ValueError(
+                "organism_ontology_term_id not found in adata.obs, you need to add an ontology term id for the organism of your anndata"
+            )
+        if not adata[0].var.index.str.contains("ENS").any() and not self.is_symbol:
+            raise ValueError(
+                "gene names in the `var.index` field of your anndata should map to the ensembl_gene nomenclature else set `is_symbol` to True if using hugo symbols"
+            )
         if adata[0].obs.organism_ontology_term_id.iloc[0] not in self.organisms:
             raise ValueError(
                 "we cannot work with this organism",

scdataloader/utils.py

@@ -154,7 +154,7 @@ def getBiomartTable(
     return res
 
 
-def validate(adata: AnnData, organism: str, need_all=True):
+def validate(adata: AnnData, organism: str, need_all=False):
     """
     validate checks if the adata object is valid for lamindb
 

{scdataloader-1.8.0.dist-info → scdataloader-1.8.1.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: scdataloader
-Version: 1.8.0
+Version: 1.8.1
 Summary: a dataloader for single cell data in lamindb
 Project-URL: repository, https://github.com/jkobject/scDataLoader
 Author-email: jkobject <jkobject@gmail.com>
@@ -16,7 +16,7 @@ Requires-Dist: harmonypy>=0.0.10
 Requires-Dist: ipykernel>=6.20.0
 Requires-Dist: lamindb[bionty,cellregistry,jupyter,ourprojects,zarr]<2,>=1.0.4
 Requires-Dist: leidenalg>=0.8.0
-Requires-Dist: lightning>=2.0.0
+Requires-Dist: lightning>=2.3.0
 Requires-Dist: matplotlib>=3.5.0
 Requires-Dist: numpy==1.26.0
 Requires-Dist: palantir>=1.3.3

scdataloader-1.8.1.dist-info/RECORD

@@ -0,0 +1,16 @@
+scdataloader/VERSION,sha256=Jc7Jc50yGOSKzF2MPUMz4dYkEhxberO83ccdD6ATS4M,6
+scdataloader/__init__.py,sha256=1SyT5MzcFl8mfp5NB4idgYQ4insXbDRd-EBNvoz_dXQ,225
+scdataloader/__main__.py,sha256=3aZnqYrH8XDT9nW9Dbb3o9kr-sx1STmXDQHxBo_h_q0,8719
+scdataloader/base.py,sha256=M1gD59OffRdLOgS1vHKygOomUoAMuzjpRtAfM3SBKF8,338
+scdataloader/collator.py,sha256=UWyTSFEYCAVcBRreFItzDgTyBx224u-ThjjW9x-osHY,12301
+scdataloader/config.py,sha256=tu9hkUiU2HfaIiVzdmrjbzt73yV4zP-t8lDuJqyGcDA,6546
+scdataloader/data.py,sha256=xWlNU6cJmrzP4BFMsJDIksLaxe1pUfgDBlQ_IeLIXj0,15578
+scdataloader/datamodule.py,sha256=6Oby-BySXaWYr34PocgCq25FLH1QUX-EsWOZI6EVjgw,21128
+scdataloader/mapped.py,sha256=DzryqhELXo-s5RgdmRFaa8zLiGjyjFKn7wW77lGLTaI,26900
+scdataloader/preprocess.py,sha256=Ewla5GYD_8YBqCDr7kaOwrYN_ok0YmYvYpwbxTComXg,35764
+scdataloader/utils.py,sha256=F5ZhdalHbxdZOs9aZ9RP9LTHGsmuoofgC39W9GS7EA4,28362
+scdataloader-1.8.1.dist-info/METADATA,sha256=NLNmj2mWRQFpwUpMxaTHuWK309MYBiBzxBxk9Nd0KD8,9946
+scdataloader-1.8.1.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
+scdataloader-1.8.1.dist-info/entry_points.txt,sha256=VXAN1m_CjbdLJ6SKYR0sBLGDV4wvv31ri7fWWuwbpno,60
+scdataloader-1.8.1.dist-info/licenses/LICENSE,sha256=OXLcl0T2SZ8Pmy2_dmlvKuetivmyPd5m1q-Gyd-zaYY,35149
+scdataloader-1.8.1.dist-info/RECORD,,

scdataloader-1.8.0.dist-info/RECORD

@@ -1,16 +0,0 @@
-scdataloader/VERSION,sha256=PrHvlLWJDKcnFYsQYUJoXIczsKzlvLTPPwrBT58GQ_Q,6
-scdataloader/__init__.py,sha256=GYqFXVzcgkqwcWodyHQSa3bnCuWsBt9jWYHEcLnx6xU,170
-scdataloader/__main__.py,sha256=3aZnqYrH8XDT9nW9Dbb3o9kr-sx1STmXDQHxBo_h_q0,8719
-scdataloader/base.py,sha256=M1gD59OffRdLOgS1vHKygOomUoAMuzjpRtAfM3SBKF8,338
-scdataloader/collator.py,sha256=n_DI630Eqo-C_G02krFD-Ixj3EKReZfW84VZy5wZHCw,11758
-scdataloader/config.py,sha256=tu9hkUiU2HfaIiVzdmrjbzt73yV4zP-t8lDuJqyGcDA,6546
-scdataloader/data.py,sha256=nLw0yCe0Sj0RGR9ioYKszwzuah-KRG0tpyjOh8xjNuY,15430
-scdataloader/datamodule.py,sha256=7xTaa6I2Yj6ikGy-bLmrsr0-9VrQUO9vW17bqhhcyJU,20972
-scdataloader/mapped.py,sha256=GCAygW7-JcEQ7sB-dsiA_nTPaA3Df5AcSd79_GFhh9k,26053
-scdataloader/preprocess.py,sha256=cHKUkGJVpnWfAVsSpl_B_IOmh8aQ0WAF2QPclhkA2eA,34876
-scdataloader/utils.py,sha256=GoRSEZ8aqmB8KussSTb95BxUBWlcLtErB_HGe0iZwic,28361
-scdataloader-1.8.0.dist-info/METADATA,sha256=UkC5E9nEXo1qf3QmNc1mz8Lvk5HL3sfmt8WiiAIqtGo,9946
-scdataloader-1.8.0.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
-scdataloader-1.8.0.dist-info/entry_points.txt,sha256=VXAN1m_CjbdLJ6SKYR0sBLGDV4wvv31ri7fWWuwbpno,60
-scdataloader-1.8.0.dist-info/licenses/LICENSE,sha256=OXLcl0T2SZ8Pmy2_dmlvKuetivmyPd5m1q-Gyd-zaYY,35149
-scdataloader-1.8.0.dist-info/RECORD,,