Add function for summing of Ensembl IDs

.gitattributes

@@ -14,6 +14,7 @@
 *.ot filter=lfs diff=lfs merge=lfs -text
 *.parquet filter=lfs diff=lfs merge=lfs -text
 *.pb filter=lfs diff=lfs merge=lfs -text
+*.pkl filter=lfs diff=lfs merge=lfs -text
 *.pt filter=lfs diff=lfs merge=lfs -text
 *.pth filter=lfs diff=lfs merge=lfs -text
 *.rar filter=lfs diff=lfs merge=lfs -text
@@ -25,5 +26,4 @@ saved_model/**/* filter=lfs diff=lfs merge=lfs -text
 *.zip filter=lfs diff=lfs merge=lfs -text
 *.zstandard filter=lfs diff=lfs merge=lfs -text
 *tfevents* filter=lfs diff=lfs merge=lfs -text
-geneformer/gene_name_id_dict.pkl filter=lfs diff=lfs merge=lfs -text
-model.safetensors filter=lfs diff=lfs merge=lfs -text
+model.safetensors filter=lfs diff=lfs merge=lfs -text

geneformer/__init__.py

@@ -4,6 +4,7 @@ from pathlib import Path
 GENE_MEDIAN_FILE = Path(__file__).parent / "gene_median_dictionary.pkl"
 TOKEN_DICTIONARY_FILE = Path(__file__).parent / "token_dictionary.pkl"
 ENSEMBL_DICTIONARY_FILE = Path(__file__).parent / "gene_name_id_dict.pkl"
+ENSEMBL_MAPPING_FILE = Path(__file__).parent / "ensembl_mapping_dict.pkl"
 
 from . import (
     collator_for_classification,

geneformer/ensembl_mapping_dict_gc95M.pkl

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:0819bcbd869cfa14279449b037eb9ed1d09a91310e77bd1a19d927465030e95c
+size 3957652

geneformer/gene_median_dictionary_gc95M.pkl

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:a51c53f6a771d64508dfaf61529df70e394c53bd20856926117ae5d641a24bf5
+size 1512661

geneformer/gene_name_id_dict_gc95M.pkl

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:fabfa0c2f49c598c59ae432a32c3499a5908c033756c663b5e0cddf58deea8e1
+size 1660882

geneformer/tokenizer.py

@@ -1,49 +1,40 @@
 """
 Geneformer tokenizer.
-
 **Input data:**
-
 | *Required format:* raw counts scRNAseq data without feature selection as .loom or anndata file.
 | *Required row (gene) attribute:* "ensembl_id"; Ensembl ID for each gene.
 | *Required col (cell) attribute:* "n_counts"; total read counts in that cell.
-
 | *Optional col (cell) attribute:* "filter_pass"; binary indicator of whether cell should be tokenized based on user-defined filtering criteria.
 | *Optional col (cell) attributes:* any other cell metadata can be passed on to the tokenized dataset as a custom attribute dictionary as shown below.
-
 **Usage:**
-
 .. code-block :: python
-
     >>> from geneformer import TranscriptomeTokenizer
     >>> tk = TranscriptomeTokenizer({"cell_type": "cell_type", "organ_major": "organ"}, nproc=4)
     >>> tk.tokenize_data("data_directory", "output_directory", "output_prefix")
-
 **Description:**
-
 | Input data is a directory with .loom or .h5ad files containing raw counts from single cell RNAseq data, including all genes detected in the transcriptome without feature selection. The input file type is specified by the argument file_format in the tokenize_data function.
-
 | The discussion below references the .loom file format, but the analagous labels are required for .h5ad files, just that they will be column instead of row attributes and vice versa due to the transposed format of the two file types.
-
 | Genes should be labeled with Ensembl IDs (loom row attribute "ensembl_id"), which provide a unique identifer for conversion to tokens. Other forms of gene annotations (e.g. gene names) can be converted to Ensembl IDs via Ensembl Biomart. Cells should be labeled with the total read count in the cell (loom column attribute "n_counts") to be used for normalization.
-
 | No cell metadata is required, but custom cell attributes may be passed onto the tokenized dataset by providing a dictionary of custom attributes to be added, which is formatted as loom_col_attr_name : desired_dataset_col_attr_name. For example, if the original .loom dataset has column attributes "cell_type" and "organ_major" and one would like to retain these attributes as labels in the tokenized dataset with the new names "cell_type" and "organ", respectively, the following custom attribute dictionary should be provided: {"cell_type": "cell_type", "organ_major": "organ"}.
-
 | Additionally, if the original .loom file contains a cell column attribute called "filter_pass", this column will be used as a binary indicator of whether to include these cells in the tokenized data. All cells with "1" in this attribute will be tokenized, whereas the others will be excluded. One may use this column to indicate QC filtering or other criteria for selection for inclusion in the final tokenized dataset.
-
 | If one's data is in other formats besides .loom or .h5ad, one can use the relevant tools (such as Anndata tools) to convert the file to a .loom or .h5ad format prior to running the transcriptome tokenizer.
-
 """
 
 from __future__ import annotations
 
+import os
 import logging
 import pickle
 import warnings
 from pathlib import Path
 from typing import Literal
+from tqdm import tqdm
+from collections import Counter
 
-import anndata as ad
 import numpy as np
+import scanpy as sc
+import loompy as lp
+import pandas as pd
 import scipy.sparse as sp
 from datasets import Dataset
 
@@ -52,7 +43,7 @@ import loompy as lp  # noqa
 
 logger = logging.getLogger(__name__)
 
-from . import GENE_MEDIAN_FILE, TOKEN_DICTIONARY_FILE
+from . import GENE_MEDIAN_FILE, TOKEN_DICTIONARY_FILE, ENSEMBL_MAPPING_FILE
 
 
 def rank_genes(gene_vector, gene_tokens):
@@ -74,6 +65,134 @@ def tokenize_cell(gene_vector, gene_tokens):
     # rank by median-scaled gene values
     return rank_genes(gene_vector[nonzero_mask], gene_tokens[nonzero_mask])
 
+def sum_ensembl_ids(data_directory,
+                    collapse_gene_ids,
+                    gene_mapping_dict,
+                    gene_token_dict,
+                    file_format = "loom",
+                    chunk_size = 512):
+
+    if file_format == "loom":
+        """ 
+        Map Ensembl IDs from gene mapping dictionary. If duplicate Ensembl IDs are found, sum counts together.
+        """
+        with lp.connect(data_directory) as data:
+            assert "ensembl_id" in data.ra.keys(), "'ensembl_id' column missing from data.ra.keys()"
+            gene_ids_in_dict = [gene for gene in data.ra.ensembl_id if gene in gene_token_dict.keys()]
+            if len(gene_ids_in_dict) == len(set(gene_ids_in_dict)):
+                token_genes_unique = True
+            else:
+                token_genes_unique = False
+            if collapse_gene_ids is False:
+                if token_genes_unique:
+                    return data_directory
+                else:
+                    raise ValueError("Error: data Ensembl IDs non-unique.")
+
+            gene_ids_collapsed = [gene_mapping_dict.get(gene_id.upper()) for gene_id in data.ra.ensembl_id]
+            gene_ids_collapsed_in_dict = [gene for gene in gene_ids_collapsed if gene in gene_token_dict.keys()]
+
+            if (len(set(gene_ids_collapsed_in_dict)) == len(set(gene_ids_in_dict))) and token_genes_unique:
+                return data_directory
+            else:
+                dedup_filename = data_directory.with_name(data_directory.stem + "__dedup.loom")
+                data.ra["gene_ids_collapsed"] = gene_ids_collapsed
+                dup_genes = [idx for idx, count in Counter(data.ra["gene_ids_collapsed"]).items() if count > 1]
+                num_chunks = int(np.ceil(data.shape[1] / chunk_size))
+                first_chunk = True
+                for _, _, view in tqdm(data.scan(axis = 1, batch_size = chunk_size), total = num_chunks):
+                    def process_chunk(view, duplic_genes):
+                        data_count_view = pd.DataFrame(view, index=data.ra["gene_ids_collapsed"])
+                        unique_data_df = data_count_view.loc[~data_count_view.index.isin(duplic_genes)]
+                        dup_data_df = data_count_view.loc[data_count_view.index.isin([i for i in duplic_genes if "None" not in i])]
+                        summed_data = dup_data_df.groupby(dup_data_df.index).sum()
+                        if not summed_data.index.is_unique:
+                            raise ValueError("Error: Ensembl IDs in summed data frame non-unique.")
+                        data_count_view = pd.concat([unique_data_df, summed_data], axis=0)
+                        if not data_count_view.index.is_unique:
+                            raise ValueError("Error: Ensembl IDs in final data frame non-unique.")
+                        return data_count_view
+                    processed_chunk = process_chunk(view[:, :], dup_genes)
+                    processed_array = processed_chunk.to_numpy()
+                    new_row_attrs = {"ensembl_id": processed_chunk.index.to_numpy()}
+
+                    if "n_counts" not in view.ca.keys():
+                        total_count_view = np.sum(view[:,:], axis=0).astype(int)
+                        view.ca["n_counts"] = total_count_view
+
+                    if first_chunk: # Create the Loom file with the first chunk
+                        lp.create(f"{dedup_filename}", processed_array, row_attrs=new_row_attrs, col_attrs=view.ca)
+                        first_chunk = False
+                    else: # Append subsequent chunks
+                        with lp.connect(dedup_filename, mode='r+') as dsout:
+                            dsout.add_columns(processed_array, col_attrs=view.ca)
+                return dedup_filename
+                
+    elif file_format == "h5ad":
+        """
+        Map Ensembl IDs from gene mapping dictionary. If duplicate Ensembl IDs are found, sum counts together.
+        Returns adata object with deduplicated Ensembl IDs.
+        """
+
+        data = sc.read_h5ad(str(data_directory))
+
+        assert "ensembl_id" in data.var.columns, "'ensembl_id' column missing from data.var"
+        gene_ids_in_dict = [gene for gene in data.var.ensembl_id if gene in gene_token_dict.keys()]
+        if len(gene_ids_in_dict) == len(set(gene_ids_in_dict)):
+            token_genes_unique = True
+        else:
+            token_genes_unique = False
+        if collapse_gene_ids is False:
+            if token_genes_unique:
+                return data
+            else:
+                raise ValueError("Error: data Ensembl IDs non-unique.")
+
+        gene_ids_collapsed = [gene_mapping_dict.get(gene_id.upper()) for gene_id in data.var.ensembl_id]
+        gene_ids_collapsed_in_dict = [gene for gene in gene_ids_collapsed if gene in gene_token_dict.keys()]
+        if (len(set(gene_ids_collapsed_in_dict)) == len(set(gene_ids_in_dict))) and token_genes_unique:
+            return data
+        
+        else:
+            data.var["gene_ids_collapsed"] = gene_ids_collapsed
+            data.var_names = gene_ids_collapsed
+            data = data[:, ~data.var.index.isna()]
+            dup_genes = [idx for idx, count in Counter(data.var_names).items() if count > 1]
+
+            num_chunks = int(np.ceil(data.shape[0] / chunk_size))
+
+            processed_genes = []
+            for i in tqdm(range(num_chunks)):
+                
+                start_idx = i * chunk_size
+                end_idx = min((i + 1) * chunk_size, data.shape[0])
+                data_chunk = data[start_idx:end_idx, :]
+
+                processed_chunks = []
+                for dup_gene in dup_genes:
+                    data_dup_gene = data_chunk[:, data_chunk.var_names == dup_gene]
+                    df = pd.DataFrame.sparse.from_spmatrix(data_dup_gene.X,
+                                                           index=data_dup_gene.obs_names,
+                                                           columns=data_dup_gene.var_names)
+                    df_sum = pd.DataFrame(df.sum(axis=1))
+                    df_sum.columns = [dup_gene] 
+                    df_sum.index = data_dup_gene.obs.index
+                    processed_chunks.append(df_sum)
+
+                processed_chunks = pd.concat(processed_chunks, axis=1) 
+                processed_genes.append(processed_chunks)
+            processed_genes = pd.concat(processed_genes, axis = 0)
+            var_df = pd.DataFrame({"gene_ids_collapsed" : processed_genes.columns})
+            var_df.index = processed_genes.columns
+            processed_genes = sc.AnnData(X = processed_genes,
+                                        obs = data.obs,
+                                        var = var_df)
+
+            data_dedup = data[:, ~data.var.index.isin(dup_genes)] # Deduplicated data
+            data_dedup = sc.concat([data_dedup, processed_genes], axis = 1)
+            data_dedup.obs = data.obs
+            data_dedup.var = data_dedup.var.rename(columns = {"gene_ids_collapsed" : "ensembl_id"})
+            return data_dedup
 
 class TranscriptomeTokenizer:
     def __init__(
@@ -83,14 +202,14 @@ class TranscriptomeTokenizer:
         chunk_size=512,
         model_input_size=2048,
         special_token=False,
+        collapse_gene_ids=True,
         gene_median_file=GENE_MEDIAN_FILE,
         token_dictionary_file=TOKEN_DICTIONARY_FILE,
+        gene_mapping_file=ENSEMBL_MAPPING_FILE,
     ):
         """
         Initialize tokenizer.
-
         **Parameters:**
-
         custom_attr_name_dict : None, dict
             | Dictionary of custom attributes to be added to the dataset.
             | Keys are the names of the attributes in the loom file.
@@ -103,12 +222,15 @@ class TranscriptomeTokenizer:
             | Max input size of model to truncate input to.
         special_token : bool = False
             | Adds CLS token before and EOS token after rank value encoding.
+        collapse_gene_ids : bool = True
+            | Whether to collapse gene IDs based on gene mapping dictionary.
         gene_median_file : Path
             | Path to pickle file containing dictionary of non-zero median
             | gene expression values across Genecorpus-30M.
         token_dictionary_file : Path
             | Path to pickle file containing token dictionary (Ensembl IDs:token).
-
+        gene_mapping_file : None, Path
+            | Path to pickle file containing dictionary for collapsing gene IDs.
         """
         # dictionary of custom attributes {output dataset column name: input .loom column name}
         self.custom_attr_name_dict = custom_attr_name_dict
@@ -134,9 +256,31 @@ class TranscriptomeTokenizer:
         with open(token_dictionary_file, "rb") as f:
             self.gene_token_dict = pickle.load(f)
 
+        # check for special token in gene_token_dict
+        if self.special_token:
+            if ("<cls>" not in self.gene_token_dict.keys()) and ("<eos>" not in self.gene_token_dict.keys()):
+                logger.error(
+                            "<cls> and <eos> required in gene_token_dict when special_token = True."
+                        )
+                raise
+
+        # if collapsing duplicate gene IDs
+        self.collapse_gene_ids = collapse_gene_ids
+
+        # load gene mappings dictionary (Ensembl IDs:Ensembl ID)
+        if gene_mapping_file is not None:
+            with open(gene_mapping_file, "rb") as f:
+                self.gene_mapping_dict = pickle.load(f)
+        else:
+            self.gene_mapping_dict = {k:k for k,_ in self.gene_token_dict.items()}
+
         # gene keys for full vocabulary
         self.gene_keys = list(self.gene_token_dict.keys())
 
+        #  Filter gene mapping dict for items that exist in gene_token_dict
+        gene_keys_set = set(self.gene_token_dict.keys())
+        self.gene_mapping_dict = {k: v for k, v in self.gene_mapping_dict.items() if v in gene_keys_set}
+
         # protein-coding and miRNA gene list dictionary for selecting .loom rows for tokenization
         self.genelist_dict = dict(zip(self.gene_keys, [True] * len(self.gene_keys)))
 
@@ -150,9 +294,7 @@ class TranscriptomeTokenizer:
     ):
         """
         Tokenize .loom files in data_directory and save as tokenized .dataset in output_directory.
-
         **Parameters:**
-
         data_directory : Path
             | Path to directory containing loom files or anndata files
         output_directory : Path
@@ -163,7 +305,6 @@ class TranscriptomeTokenizer:
             | Format of input files. Can be "loom" or "h5ad".
         use_generator : bool
             | Whether to use generator or dict for tokenization.
-
         """
         tokenized_cells, cell_metadata = self.tokenize_files(
             Path(data_directory), file_format
@@ -214,7 +355,7 @@ class TranscriptomeTokenizer:
         return tokenized_cells, cell_metadata
 
     def tokenize_anndata(self, adata_file_path, target_sum=10_000):
-        adata = ad.read(adata_file_path, backed="r")
+        adata = sum_ensembl_ids(adata_file_path, self.collapse_gene_ids, self.gene_mapping_dict, self.gene_token_dict, file_format = "h5ad", chunk_size = self.chunk_size)
 
         if self.custom_attr_name_dict is not None:
             file_cell_metadata = {
@@ -256,7 +397,8 @@ class TranscriptomeTokenizer:
             idx = filter_pass_loc[i : i + self.chunk_size]
 
             n_counts = adata[idx].obs["n_counts"].values[:, None]
-            X_view = adata[idx, coding_miRNA_loc].X
+            X_view0 = adata[idx,:].X
+            X_view = X_view0[:, coding_miRNA_loc]
             X_norm = X_view / n_counts * target_sum / norm_factor_vector
             X_norm = sp.csr_matrix(X_norm)
 
@@ -280,6 +422,9 @@ class TranscriptomeTokenizer:
                 attr_key: [] for attr_key in self.custom_attr_name_dict.keys()
             }
 
+        dedup_filename = loom_file_path.with_name(loom_file_path.stem + "__dedup.loom")
+        loom_file_path = sum_ensembl_ids(loom_file_path, self.collapse_gene_ids, self.gene_mapping_dict, self.gene_token_dict, file_format = "loom", chunk_size = self.chunk_size)
+
         with lp.connect(str(loom_file_path)) as data:
             # define coordinates of detected protein-coding or miRNA genes and vector of their normalization factors
             coding_miRNA_loc = np.where(
@@ -341,6 +486,9 @@ class TranscriptomeTokenizer:
                 else:
                     file_cell_metadata = None
 
+        if str(dedup_filename) == str(loom_file_path):
+            os.remove(str(dedup_filename))
+
         return tokenized_cells, file_cell_metadata
 
     def create_dataset(
@@ -396,4 +544,4 @@ class TranscriptomeTokenizer:
         output_dataset_truncated = output_dataset.map(
             format_cell_features, num_proc=self.nproc
         )
-        return output_dataset_truncated
+        return output_dataset_truncated