--- a/decoupler_aucell_score.py	Tue Apr 16 11:49:19 2024 +0000
+++ b/decoupler_aucell_score.py	Mon Jul 15 10:56:37 2024 +0000
@@ -1,16 +1,15 @@
 import argparse
-import os
-import tempfile
 
 import anndata
 import decoupler as dc
+import numba as nb
 import pandas as pd
-import numba as nb
 
 
 def read_gmt_long(gmt_file):
-    """
-    Reads a GMT file and produce a Pandas DataFrame in long format, ready to be passed to the AUCell method.
+    r"""
+    Reads a GMT file and produce a Pandas DataFrame in long format, ready to
+    be passed to the AUCell method.
 
     Parameters
     ----------
@@ -20,9 +19,29 @@
     Returns
     -------
     pd.DataFrame
-        A DataFrame with the gene sets. Each row represents a gene set to gene assignment, and the columns are "gene_set_name" and "genes".
-    >>> line = "HALLMARK_NOTCH_SIGNALING\\thttp://www.gsea-msigdb.org/gsea/msigdb/human/geneset/HALLMARK_NOTCH_SIGNALING\\tJAG1\\tNOTCH3\\tNOTCH2\\tAPH1A\\tHES1\\tCCND1\\tFZD1\\tPSEN2\\tFZD7\\tDTX1\\tDLL1\\tFZD5\\tMAML2\\tNOTCH1\\tPSENEN\\tWNT5A\\tCUL1\\tWNT2\\tDTX4\\tSAP30\\tPPARD\\tKAT2A\\tHEYL\\tSKP1\\tRBX1\\tTCF7L2\\tARRB1\\tLFNG\\tPRKCA\\tDTX2\\tST3GAL6\\tFBXW11\\n"
-    >>> line2 = "HALLMARK_APICAL_SURFACE\\thttp://www.gsea-msigdb.org/gsea/msigdb/human/geneset/HALLMARK_APICAL_SURFACE\\tB4GALT1\\tRHCG\\tMAL\\tLYPD3\\tPKHD1\\tATP6V0A4\\tCRYBG1\\tSHROOM2\\tSRPX\\tMDGA1\\tTMEM8B\\tTHY1\\tPCSK9\\tEPHB4\\tDCBLD2\\tGHRL\\tLYN\\tGAS1\\tFLOT2\\tPLAUR\\tAKAP7\\tATP8B1\\tEFNA5\\tSLC34A3\\tAPP\\tGSTM3\\tHSPB1\\tSLC2A4\\tIL2RB\\tRTN4RL1\\tNCOA6\\tSULF2\\tADAM10\\tBRCA1\\tGATA3\\tAFAP1L2\\tIL2RG\\tCD160\\tADIPOR2\\tSLC22A12\\tNTNG1\\tSCUBE1\\tCX3CL1\\tCROCC\\n"
+        A DataFrame with the gene sets. Each row represents a gene set to gene
+        assignment, and the columns are "gene_set_name" and "genes".
+    >>> import os
+    >>> import tempfile
+    >>> line = "HALLMARK_NOTCH_SIGNALING\
+    ... \thttp://www.gsea-msigdb.org/\
+    ... gsea/msigdb/human/geneset/HALLMARK_NOTCH_SIGNALING\
+    ... \tJAG1\tNOTCH3\tNOTCH2\tAPH1A\tHES1\tCCND1\
+    ... \tFZD1\tPSEN2\tFZD7\tDTX1\tDLL1\tFZD5\tMAML2\
+    ... \tNOTCH1\tPSENEN\tWNT5A\tCUL1\tWNT2\tDTX4\
+    ... \tSAP30\tPPARD\tKAT2A\tHEYL\tSKP1\tRBX1\tTCF7L2\
+    ... \tARRB1\tLFNG\tPRKCA\tDTX2\tST3GAL6\tFBXW11\n"
+    >>> line2 = "HALLMARK_APICAL_SURFACE\
+    ... \thttp://www.gsea-msigdb.org/\
+    ... gsea/msigdb/human/geneset/HALLMARK_APICAL_SURFACE\
+    ... \tB4GALT1\tRHCG\tMAL\tLYPD3\tPKHD1\tATP6V0A4\
+    ... \tCRYBG1\tSHROOM2\tSRPX\tMDGA1\tTMEM8B\tTHY1\
+    ... \tPCSK9\tEPHB4\tDCBLD2\tGHRL\tLYN\tGAS1\tFLOT2\
+    ... \tPLAUR\tAKAP7\tATP8B1\tEFNA5\tSLC34A3\tAPP\
+    ... \tGSTM3\tHSPB1\tSLC2A4\tIL2RB\tRTN4RL1\tNCOA6\
+    ... \tSULF2\tADAM10\tBRCA1\tGATA3\tAFAP1L2\tIL2RG\
+    ... \tCD160\tADIPOR2\tSLC22A12\tNTNG1\tSCUBE1\tCX3CL1\
+    ... \tCROCC\n"
     >>> temp_dir = tempfile.gettempdir()
     >>> temp_gmt = os.path.join(temp_dir, "temp_file.gmt")
     >>> with open(temp_gmt, "w") as f:
@@ -36,7 +55,8 @@
     >>> len(df.loc[df["gene_set"] == "HALLMARK_APICAL_SURFACE"].gene.tolist())
     44
     """
-    # Create a list of dictionaries, where each dictionary represents a gene set
+    # Create a list of dictionaries, where each dictionary represents a
+    # gene set
     gene_sets = {}
 
     # Read the GMT file into a list of lines
@@ -46,12 +66,20 @@
             if not line:
                 break
             fields = line.strip().split("\t")
-            gene_sets[fields[0]]= fields[2:]
+            gene_sets[fields[0]] = fields[2:]
 
-    return pd.concat(pd.DataFrame({'gene_set':k, 'gene':v}) for k, v in gene_sets.items())
+    return pd.concat(
+        pd.DataFrame({"gene_set": k, "gene": v}) for k, v in gene_sets.items()
+    )
 
 
-def score_genes_aucell_mt(adata: anndata.AnnData, gene_set_gene: pd.DataFrame, use_raw=False, min_n_genes=5, var_gene_symbols_field=None):
+def score_genes_aucell_mt(
+    adata: anndata.AnnData,
+    gene_set_gene: pd.DataFrame,
+    use_raw=False,
+    min_n_genes=5,
+    var_gene_symbols_field=None,
+):
     """Score genes using Aucell.
 
     Parameters
@@ -60,17 +88,23 @@
     gene_set_gene: pd.DataFrame with columns gene_set and gene
     use_raw : bool, optional, False by default.
     min_n_genes : int, optional, 5 by default.
-    var_gene_symbols_field : str, optional, None by default. The field in var where gene symbols are stored
+    var_gene_symbols_field : str, optional, None by default. The field in var
+    where gene symbols are stored
 
     >>> import scanpy as sc
     >>> import decoupler as dc
     >>> adata = sc.datasets.pbmc68k_reduced()
-    >>> r_gene_list = adata.var[adata.var.index.str.startswith("RP")].index.tolist()
-    >>> m_gene_list = adata.var[adata.var.index.str.startswith("M")].index.tolist()
+    >>> r_gene_list = adata.var[
+    ...                  adata.var.index.str.startswith("RP")].index.tolist()
+    >>> m_gene_list = adata.var[
+    ...                  adata.var.index.str.startswith("M")].index.tolist()
     >>> gene_set = {}
     >>> gene_set["m"] = m_gene_list
     >>> gene_set["r"] = r_gene_list
-    >>> gene_set_df = pd.concat(pd.DataFrame({'gene_set':k, 'gene':v}) for k, v in gene_set.items())
+    >>> gene_set_df = pd.concat(
+    ...                  pd.DataFrame(
+    ...                     {'gene_set':k, 'gene':v}
+    ...     ) for k, v in gene_set.items())
     >>> score_genes_aucell_mt(adata, gene_set_df, use_raw=False)
     >>> "AUCell_m" in adata.obs.columns
     True
@@ -78,47 +112,72 @@
     True
     """
 
-    # if var_gene_symbols_fiels is provided, transform gene_set_gene df so that gene contains gene ids instead of gene symbols
+    # if var_gene_symbols_fiels is provided, transform gene_set_gene df so
+    #  that gene contains gene ids instead of gene symbols
     if var_gene_symbols_field:
-        # merge the index of var to gene_set_gene df based on var_gene_symbols_field
+        # merge the index of var to gene_set_gene df based on
+        # var_gene_symbols_field
         var_id_symbols = adata.var[[var_gene_symbols_field]]
-        var_id_symbols['gene_id'] = var_id_symbols.index
+        var_id_symbols["gene_id"] = var_id_symbols.index
 
-        gene_set_gene = gene_set_gene.merge(var_id_symbols, left_on='gene', right_on=var_gene_symbols_field, how='left')
-        # this will still produce some empty gene_ids (genes in the gene_set_gene df that are not in the var df), fill those
-        # with the original gene symbol from the gene_set to avoid deforming the AUCell calculation
-        gene_set_gene['gene_id'] = gene_set_gene['gene_id'].fillna(gene_set_gene['gene'])
-        gene_set_gene['gene'] = gene_set_gene['gene_id']
-    
+        gene_set_gene = gene_set_gene.merge(
+            var_id_symbols,
+            left_on="gene",
+            right_on=var_gene_symbols_field,
+            how="left",
+        )
+        # this will still produce some empty gene_ids (genes in the
+        # gene_set_gene df that are not in the var df), fill those
+        # with the original gene symbol from the gene_set to avoid
+        # deforming the AUCell calculation
+        gene_set_gene["gene_id"] = gene_set_gene["gene_id"].fillna(
+            gene_set_gene["gene"]
+        )
+        gene_set_gene["gene"] = gene_set_gene["gene_id"]
+
     # run decoupler's run_aucell
     dc.run_aucell(
-            adata, net=gene_set_gene, source="gene_set", target="gene", use_raw=use_raw, min_n=min_n_genes
-        )
+        adata,
+        net=gene_set_gene,
+        source="gene_set",
+        target="gene",
+        use_raw=use_raw,
+        min_n=min_n_genes,
+    )
     for gs in gene_set_gene.gene_set.unique():
-        if gs in adata.obsm['aucell_estimate'].keys():
+        if gs in adata.obsm["aucell_estimate"].keys():
             adata.obs[f"AUCell_{gs}"] = adata.obsm["aucell_estimate"][gs]
 
 
 def run_for_genelists(
-    adata, gene_lists, score_names, use_raw=False, gene_symbols_field=None, min_n_genes=5
+    adata,
+    gene_lists,
+    score_names,
+    use_raw=False,
+    gene_symbols_field=None,
+    min_n_genes=5,
 ):
     if len(gene_lists) == len(score_names):
         for gene_list, score_names in zip(gene_lists, score_names):
             genes = gene_list.split(",")
             gene_sets = {}
             gene_sets[score_names] = genes
-            gene_set_gene_df = pd.concat(pd.DataFrame({'gene_set':k, 'gene':v}) for k, v in gene_sets.items())
-            
+            gene_set_gene_df = pd.concat(
+                pd.DataFrame({"gene_set": k, "gene": v})
+                for k, v in gene_sets.items()
+            )
+
             score_genes_aucell_mt(
                 adata,
                 gene_set_gene_df,
                 use_raw,
                 min_n_genes,
-                var_gene_symbols_field=gene_symbols_field
+                var_gene_symbols_field=gene_symbols_field,
             )
     else:
         raise ValueError(
-            "The number of gene lists (separated by :) and score names (separated by :) must be the same"
+            "The number of gene lists (separated by :) and score names \
+                (separated by :) must be the same"
         )
 
 
@@ -126,32 +185,41 @@
     # Create command-line arguments parser
     parser = argparse.ArgumentParser(description="Score genes using Aucell")
     parser.add_argument(
-        "--input_file", type=str, help="Path to input AnnData file", required=True
+        "--input_file",
+        type=str,
+        help="Path to input AnnData file",
+        required=True,
     )
     parser.add_argument(
         "--output_file", type=str, help="Path to output file", required=True
     )
-    parser.add_argument("--gmt_file", type=str, help="Path to GMT file", required=False)
+    parser.add_argument(
+        "--gmt_file", type=str, help="Path to GMT file", required=False
+    )
     # add argument for gene sets to score
     parser.add_argument(
         "--gene_sets_to_score",
         type=str,
         required=False,
-        help="Optional comma separated list of gene sets to score (the need to be in the gmt file)",
+        help="Optional comma separated list of gene sets to score \
+            (the need to be in the gmt file)",
     )
     # add argument for gene list (comma separated) to score
     parser.add_argument(
         "--gene_lists_to_score",
         type=str,
         required=False,
-        help="Comma separated list of genes to score. You can have more than one set of genes, separated by colon :",
+        help="Comma separated list of genes to score. You can have more \
+            than one set of genes, separated by colon :",
     )
     # argument for the score name when using the gene list
     parser.add_argument(
         "--score_names",
         type=str,
         required=False,
-        help="Name of the score column when using the gene list. You can have more than one set of score names, separated by colon :. It should be the same length as the number of gene lists.",
+        help="Name of the score column when using the gene list. You can \
+            have more than one set of score names, separated by colon :. \
+                It should be the same length as the number of gene lists.",
     )
     parser.add_argument(
         "--gene_symbols_field",
@@ -159,7 +227,8 @@
         help="Name of the gene symbols field in the AnnData object",
         required=True,
     )
-    # argument for min_n Minimum of targets per source. If less, sources are removed.
+    # argument for min_n Minimum of targets per source. If less, sources
+    # are removed.
     parser.add_argument(
         "--min_n",
         type=int,
@@ -169,11 +238,18 @@
     )
     parser.add_argument("--use_raw", action="store_true", help="Use raw data")
     parser.add_argument(
-        "--write_anndata", action="store_true", help="Write the modified AnnData object"
+        "--write_anndata",
+        action="store_true",
+        help="Write the modified AnnData object",
     )
     # argument for number of max concurrent processes
-    parser.add_argument("--max_threads", type=int, required=False, default=1, help="Number of max concurrent threads")
-
+    parser.add_argument(
+        "--max_threads",
+        type=int,
+        required=False,
+        default=1,
+        help="Number of max concurrent threads",
+    )
 
     # Parse command-line arguments
     args = parser.parse_args()
@@ -189,23 +265,40 @@
         msigdb = read_gmt_long(args.gmt_file)
 
         gene_sets_to_score = (
-            args.gene_sets_to_score.split(",") if args.gene_sets_to_score else []
+            args.gene_sets_to_score.split(",")
+            if args.gene_sets_to_score
+            else []
         )
         if gene_sets_to_score:
-            # we limit the GMT file read to the genesets specified in the gene_sets_to_score argument
+            # we limit the GMT file read to the genesets specified in the
+            # gene_sets_to_score argument
             msigdb = msigdb[msigdb["gene_set"].isin(gene_sets_to_score)]
-        
-        score_genes_aucell_mt(adata, msigdb, args.use_raw, args.min_n, var_gene_symbols_field=args.gene_symbols_field)
+
+        score_genes_aucell_mt(
+            adata,
+            msigdb,
+            args.use_raw,
+            args.min_n,
+            var_gene_symbols_field=args.gene_symbols_field,
+        )
     elif args.gene_lists_to_score is not None and args.score_names is not None:
         gene_lists = args.gene_lists_to_score.split(":")
         score_names = args.score_names.split(",")
         run_for_genelists(
-            adata, gene_lists, score_names, args.use_raw, args.gene_symbols_field, args.min_n
+            adata,
+            gene_lists,
+            score_names,
+            args.use_raw,
+            args.gene_symbols_field,
+            args.min_n,
         )
 
-    # Save the modified AnnData object or generate a file with cells as rows and the new score_names columns
+    # Save the modified AnnData object or generate a file with cells as rows
+    # and the new score_names columns
     if args.write_anndata:
         adata.write_h5ad(args.output_file)
     else:
-        new_columns = [col for col in adata.obs.columns if col.startswith("AUCell_")]
+        new_columns = [
+            col for col in adata.obs.columns if col.startswith("AUCell_")
+        ]
         adata.obs[new_columns].to_csv(args.output_file, sep="\t", index=True)