scdataloader 1.9.2__py3-none-any.whl → 2.0.2__py3-none-any.whl

scdataloader/__main__.py

@@ -1,5 +1,5 @@
 import argparse
-from typing import Optional, Union
+from typing import List, Optional, Union
 
 import lamindb as ln
 
@@ -149,7 +149,7 @@ def main():
     )
     preprocess_parser.add_argument(
         "--batch_keys",
-        type=list[str],
+        type=List[str],
         default=[
             "assay_ontology_term_id",
             "self_reported_ethnicity_ontology_term_id",
@@ -229,11 +229,11 @@ def main():
     if args.instance is not None:
         collection = (
             ln.Collection.using(instance=args.instance)
-            .filter(name=args.name, version=args.version)
+            .filter(key=args.name, version=args.version)
             .first()
         )
     else:
-        collection = ln.Collection.filter(name=args.name, version=args.version).first()
+        collection = ln.Collection.filter(key=args.name, version=args.version).first()
 
     print(
         "using the dataset ", collection, " of size ", len(collection.artifacts.all())
@@ -262,7 +262,6 @@ def main():
         additional_preprocess=additional_preprocess,
         additional_postprocess=additional_postprocess,
         keep_files=False,
-        force_preloaded=args.force_preloaded,
     )
 
     # Preprocess the dataset

scdataloader/collator.py

@@ -1,18 +1,20 @@
-from typing import Optional
+from typing import List, Optional
 
 import numpy as np
+import pandas as pd
 from torch import Tensor, long
 
-from .utils import downsample_profile, load_genes
+from .preprocess import _digitize
+from .utils import load_genes
 
 
 class Collator:
     def __init__(
         self,
-        organisms: list[str],
+        organisms: List[str],
         how: str = "all",
         org_to_id: dict[str, int] = None,
-        valid_genes: list[str] = [],
+        valid_genes: Optional[List[str]] = None,
         max_len: int = 2000,
         add_zero_genes: int = 0,
         logp1: bool = False,
@@ -20,10 +22,9 @@ class Collator:
         n_bins: int = 0,
         tp_name: Optional[str] = None,
         organism_name: str = "organism_ontology_term_id",
-        class_names: list[str] = [],
-        genelist: list[str] = [],
-        downsample: Optional[float] = None,  # don't use it for training!
-        save_output: Optional[str] = None,
+        class_names: List[str] = [],
+        genelist: List[str] = [],
+        genedf: Optional[pd.DataFrame] = None,
     ):
         """
         This class is responsible for collating data for the scPRINT model. It handles the
@@ -57,13 +58,8 @@ class Collator:
             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
-            downsample (float, optional): Downsample the profile to a certain number of cells. Defaults to None.
-                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
         """
         self.organisms = organisms
-        self.genedf = load_genes(organisms)
         self.max_len = max_len
         self.n_bins = n_bins
         self.add_zero_genes = add_zero_genes
@@ -75,32 +71,36 @@ class Collator:
         self.organism_name = organism_name
         self.tp_name = tp_name
         self.class_names = class_names
-        self.save_output = save_output
         self.start_idx = {}
         self.accepted_genes = {}
-        self.downsample = downsample
         self.to_subset = {}
-        self._setup(org_to_id, valid_genes, genelist)
+        self._setup(genedf, org_to_id, valid_genes, genelist)
 
-    def _setup(self, org_to_id=None, valid_genes=[], genelist=[]):
+    def _setup(self, genedf=None, org_to_id=None, valid_genes=[], genelist=[]):
+        if genedf is None:
+            genedf = load_genes(self.organisms)
+            self.organism_ids = (
+                set([org_to_id[k] for k in self.organisms])
+                if org_to_id is not None
+                else set(self.organisms)
+            )
         self.org_to_id = org_to_id
         self.to_subset = {}
         self.accepted_genes = {}
         self.start_idx = {}
-        self.organism_ids = (
-            set([org_to_id[k] for k in self.organisms])
-            if org_to_id is not None
-            else set(self.organisms)
-        )
+
+        if valid_genes is not None:
+            if len(set(valid_genes) - set(genedf.index)) > 0:
+                print("Some valid genes are not in the genedf!!!")
+            tot = genedf[genedf.index.isin(valid_genes)]
+        else:
+            tot = genedf
         for organism in self.organisms:
-            ogenedf = self.genedf[self.genedf.organism == organism]
-            if len(valid_genes) > 0:
-                tot = self.genedf[self.genedf.index.isin(valid_genes)]
-            else:
-                tot = self.genedf
             org = org_to_id[organism] if org_to_id is not None else organism
             self.start_idx.update({org: np.where(tot.organism == organism)[0][0]})
-            if len(valid_genes) > 0:
+
+            ogenedf = genedf[genedf.organism == organism]
+            if valid_genes is not None:
                 self.accepted_genes.update({org: ogenedf.index.isin(valid_genes)})
             if len(genelist) > 0:
                 df = ogenedf[ogenedf.index.isin(valid_genes)]
@@ -111,7 +111,7 @@ class Collator:
         __call__ applies the collator to a minibatch of data
 
         Args:
-            batch (list[dict[str: array]]): List of dicts of arrays containing gene expression data.
+            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
@@ -119,7 +119,7 @@ class Collator:
                 elem["class_names.."]: other classes
 
         Returns:
-            list[Tensor]: List of tensors containing the collated data.
+            List[Tensor]: List of tensors containing the collated data.
         """
         # do count selection
         # get the unseen info and don't add any unseen
@@ -133,6 +133,7 @@ class Collator:
         nnz_loc = []
         is_meta = []
         knn_cells = []
+        knn_cells_info = []
         for elem in batch:
             organism_id = elem[self.organism_name]
             if organism_id not in self.organism_ids:
@@ -188,7 +189,14 @@ class Collator:
                 if "knn_cells" in elem:
                     # we complete with genes expressed in the knn
                     # which is not a zero_loc in this context
-                    zero_loc = np.argsort(elem["knn_cells"].sum(0))[-ma:][::-1]
+                    knn_expr = elem["knn_cells"].sum(0)
+                    mask = np.ones(len(knn_expr), dtype=bool)
+                    mask[loc] = False
+                    available_indices = np.where(mask)[0]
+                    available_knn_expr = knn_expr[available_indices]
+                    sorted_indices = np.argsort(available_knn_expr)[::-1]
+                    selected = min(ma, len(available_indices))
+                    zero_loc = available_indices[sorted_indices[:selected]]
                 else:
                     zero_loc = np.where(expr == 0)[0]
                     zero_loc = zero_loc[
@@ -212,6 +220,8 @@ class Collator:
             exprs.append(expr)
             if "knn_cells" in elem:
                 knn_cells.append(elem["knn_cells"])
+            if "knn_cells_info" in elem:
+                knn_cells_info.append(elem["knn_cells_info"])
             # then we need to add the start_idx to the loc to give it the correct index
             # according to the model
             gene_locs.append(loc + self.start_idx[organism_id])
@@ -231,15 +241,46 @@ class Collator:
         dataset = np.array(dataset)
         is_meta = np.array(is_meta)
         knn_cells = np.array(knn_cells)
+        knn_cells_info = np.array(knn_cells_info)
+
         # normalize counts
         if self.norm_to is not None:
             expr = (expr * self.norm_to) / total_count[:, None]
+            # TODO: solve issue here
+            knn_cells = (knn_cells * self.norm_to) / total_count[:, None]
         if self.logp1:
             expr = np.log2(1 + expr)
+            knn_cells = np.log2(1 + knn_cells)
 
         # do binning of counts
-        if self.n_bins:
-            pass
+        if self.n_bins > 0:
+            binned_rows = []
+            bin_edges = []
+            for row in expr:
+                if row.max() == 0:
+                    print(
+                        "The input data contains all zero rows. Please make sure "
+                        "this is expected. You can use the `filter_cell_by_counts` "
+                        "arg to filter out all zero rows."
+                    )
+                    binned_rows.append(np.zeros_like(row, dtype=np.int64))
+                    bin_edges.append(np.array([0] * self.n_bins))
+                    continue
+                non_zero_ids = row.nonzero()
+                non_zero_row = row[non_zero_ids]
+                bins = np.quantile(non_zero_row, np.linspace(0, 1, self.n_bins - 1))
+                # bins = np.sort(np.unique(bins))
+                # NOTE: comment this line for now, since this will make the each category
+                # has different relative meaning across datasets
+                non_zero_digits = _digitize(non_zero_row, bins)
+                assert non_zero_digits.min() >= 1
+                assert non_zero_digits.max() <= self.n_bins - 1
+                binned_row = np.zeros_like(row, dtype=np.int64)
+                binned_row[non_zero_ids] = non_zero_digits
+                binned_rows.append(binned_row)
+                bin_edges.append(np.concatenate([[0], bins]))
+            expr = np.stack(binned_rows)
+            # expr = np.digitize(expr, bins=self.bins)
 
         ret = {
             "x": Tensor(expr),
@@ -252,51 +293,8 @@ class Collator:
             ret.update({"is_meta": Tensor(is_meta).int()})
         if len(knn_cells) > 0:
             ret.update({"knn_cells": Tensor(knn_cells)})
+        if len(knn_cells_info) > 0:
+            ret.update({"knn_cells_info": Tensor(knn_cells_info)})
         if len(dataset) > 0:
             ret.update({"dataset": Tensor(dataset).to(long)})
-        if self.downsample is not None:
-            ret["x"] = downsample_profile(ret["x"], self.downsample)
-        if self.save_output is not None:
-            with open(self.save_output, "a") as f:
-                np.savetxt(f, ret["x"].numpy())
-            with open(self.save_output + "_loc", "a") as f:
-                np.savetxt(f, gene_locs)
         return ret
-
-
-#############
-#### WIP ####
-#############
-class GeneformerCollator(Collator):
-    def __init__(self, *args, gene_norm_list: list, **kwargs):
-        """
-        GeneformerCollator to finish
-
-        Args:
-            gene_norm_list (list): the normalization of expression through all datasets, per gene.
-        """
-        super().__init__(*args, **kwargs)
-        self.gene_norm_list = gene_norm_list
-
-    def __call__(self, batch):
-        super().__call__(batch)
-        # normlization per gene
-
-        # tokenize the empty locations
-
-
-class scGPTCollator(Collator):
-    """
-    scGPTCollator to finish
-    """
-
-    def __call__(self, batch):
-        super().__call__(batch)
-        # binning
-
-        # tokenize the empty locations
-
-
-class scPRINTCollator(Collator):
-    def __call__(self, batch):
-        super().__call__(batch)

scdataloader/config.py

@@ -113,26 +113,34 @@ COARSE_ASSAY = {
 
 
 MAIN_HUMAN_MOUSE_DEV_STAGE_MAP = {
-    "HsapDv:0010000": [
+    "HsapDv:0010000": [  # postnatal stage
         "MmusDv:0000092",  # postnatal stage
     ],
-    "HsapDv:0000258": [  # mature stage
+    "HsapDv:0000258": [  # mature stage >15
         "MmusDv:0000110",  # mature stage
-        "HsapDv:0000204", # 
+        "HsapDv:0000204",  #
     ],
-    "HsapDv:0000227": [  # late adult stage
+    "HsapDv:0000087": [],  # adult stage >19
+    "HsapDv:0000227": [  # late adult stage > 40
         "MmusDv:0000091",  # 20 month-old stage
         "MmusDv:0000089",  # 18 month-old stage
+        "HsapDv:0000091",  # > 45
+        "HsapDv:0000093",  # > 65
+    ],
+    "HsapDv:0000272": [  # 60-79 year-old stage
+        "HsapDv:0000094",  # 60-79 year-old stage
     ],
-    "HsapDv:0000272": [],  # 60-79 year-old stage
     "HsapDv:0000095": [],  # 80 year-old and over stage
-    "HsapDv:0000267": [  # middle aged stage
+    "HsapDv:0000267": [  # middle aged stage >40 <60
         "MmusDv:0000087",  # 16 month-old stage
         "UBERON:0018241",  # prime adult stage
         "MmusDv:0000083",  # 12 month-old stage
         "HsapDv:0000092",  # same
     ],
-    "HsapDv:0000266": [  # young adult stage
+    "HsapDv:0000266": [  # young adult stage <40
+        "HsapDv:0000088",  # mature stage
+        "HsapDv:0000090",  # 25 - 44
+        "HsapDv:0000086",  # adolescent stage
         "MmusDv:0000050",  # 6 weeks
         "HsapDv:0000089",  # same
         "MmusDv:0000051",  # 7 weeks
@@ -163,22 +171,30 @@ MAIN_HUMAN_MOUSE_DEV_STAGE_MAP = {
         "MmusDv:0000099",  # 26 weeks
         "MmusDv:0000102",  # 29 weeks
     ],
-    "HsapDv:0000265": [],  # child stage (1-4 yo)
+    "HsapDv:0000265": [  # child stage (1-4 yo)
+        "HsapDv:0000084",  # 2-5 yo
+    ],
     "HsapDv:0000271": [  # juvenile stage (5-14 yo)
         "MmusDv:0000048",  # 4 weeks
         "MmusDv:0000049",  # 5 weeks
+        "HsapDv:0000081",  # child
+        "HsapDv:0000085",  # 6-11 yo
     ],
-    "HsapDv:0000260": [  # infant stage
+    "HsapDv:0000260": [  # infant stage <2
         "MmusDv:0000046",  # 2 weeks
         "MmusDv:0000045",  # 1 week
         "MmusDv:0000047",  # 3 weeks
         "HsapDv:0000083",
+        "HsapDv:0000256",  # under 1 yo
     ],
     "HsapDv:0000262": [  # newborn stage (0-28 days)
         "MmusDv:0000036",  # Theiler stage 27
         "MmusDv:0000037",  # Theiler stage 28
         "MmusDv:0000113",  # 4-7 days
+        "HsapDv:0000174",  # 1 month-old stage
+        "HsapDv:0000082",  # newborn stage
     ],
+    "HsapDv:0000002": [],  # embryonic stage
     "HsapDv:0000007": [],  # Carnegie stage 03
     "HsapDv:0000008": [],  # Carnegie stage 04
     "HsapDv:0000009": [],  # Carnegie stage 05