sweep removed from count.py and added to sweep.py

Author: josephrich98

kb_python/count.py

@@ -86,64 +86,6 @@
 
 INSPECT_PARSER = re.compile(r'^.*?(?P<count>[0-9]+)')
 
-def check_kwargs(kwargs):
-    if not kwargs:
-        return
-    try:
-        import inspect as inspect_module
-        import cellsweep
-        sig = inspect_module.signature(cellsweep.denoise_count_matrix.__wrapped__)
-        cellsweep_arg_names = list(sig.parameters.keys())
-        if any(param not in cellsweep_arg_names for param in kwargs):
-            invalid_params = [
-                param for param in kwargs
-                if param not in cellsweep_arg_names
-            ]
-            raise TypeError(
-                f"count() got an unexpected keyword argument(s): {', '.join(invalid_params)}"
-            )
-    except Exception as e:
-        pass
-
-def run_cellsweep(counts_dir, out_dir, threads=2, kwargs=None):
-    try:
-        import inspect as inspect_module
-        import cellsweep
-        sig = inspect_module.signature(cellsweep.denoise_count_matrix.__wrapped__)
-        cellsweep_arg_names = list(sig.parameters.keys())
-        cellsweep_kwargs = {}
-        if kwargs:
-            cellsweep_kwargs = {k: v for k, v in kwargs.items() if k in cellsweep_arg_names}
-        cellsweep_counts_dir = os.path.join(out_dir, "counts_swept")
-        cellsweep_adata_path = os.path.join(cellsweep_counts_dir, "swept_adata.h5ad")
-
-        matrix_path = os.path.join(counts_dir, "cells_x_genes.mtx")
-        barcodes_path = os.path.join(counts_dir, "cells_x_genes.barcodes.txt")
-        genes_path = os.path.join(counts_dir, "cells_x_genes.genes.names.txt")
-        adata = import_matrix_as_anndata(matrix_path, barcodes_path, genes_path)
-        # adata = cellsweep.utils.read_kb_mtx_as_adata(counts_dir)
-
-        # TODO:
-        #* 1. think of how to do automatic celltyping
-        #* 2. implement the requirement for expected_cells or umi_cutoff, or have a way to auto-detect
-
-        # add celltypes
-        # adata = cs_utils.determine_cell_types(adata, model_pkl=model_pkl, filter_empty=True, expected_cells=expected_cells, celltypist_convert=celltypist_convert, celltypist_map_file=celltypist_map_file, verbose=verbose)
-        
-        import numpy as np  #!!! erase
-        celltypes = ["celltype1", "celltype2", "celltype3"]   #!!! erase
-        adata.obs["celltype"] = np.random.choice(celltypes, size=adata.n_obs)  #!!! erase
-
-        _ = cellsweep.denoise_count_matrix(
-            adata=adata,
-            adata_out=cellsweep_adata_path,
-            threads=threads,
-            **cellsweep_kwargs
-        )
-        return cellsweep_adata_path
-    except Exception as e:
-        logger.error(f"Error running cellsweep: {e}")
-        return None
 
 def make_transcript_t2g(txnames_path: str, out_path: str) -> str:
     """Make a two-column t2g file from a transcripts file
@@ -1313,8 +1255,6 @@ def count(
     quant_umis: bool = False,
     keep_flags: bool = False,
     exact_barcodes: bool = False,
-    remove_ambient: bool = False,
-    **kwargs
 ) -> Dict[str, Union[str, Dict[str, str]]]:
     """Generates count matrices for single-cell RNA seq.
 
@@ -1392,18 +1332,13 @@ def count(
         quant_umis: Whether to run quant-tcc when there are UMIs, defaults to `False`
         keep_flags: Preserve flag column when sorting BUS file, defaults to `False`
         exact_barcodes: Use exact match for 'correcting' barcodes to on-list, defaults to `False`
-        remove_ambient: Whether to remove ambient RNA using CellSweep, defaults to `False`
-        **kwargs: Additional keyword arguments to pass to CellSweep
 
     Returns:
         Dictionary containing paths to generated files
     """
     STATS.start()
     is_batch = isinstance(fastqs, str)
 
-    #* kwargs is only added for cellsweep, so check accordingly
-    check_kwargs(kwargs)
-
     results = {}
     make_directory(out_dir)
     unfiltered_results = results.setdefault('unfiltered', {})
@@ -1820,10 +1755,6 @@ def update_results_with_suffix(current_results, new_results, suffix):
             temp_dir=temp_dir
         )
         unfiltered_results.update(report_result)
-    
-    if remove_ambient:
-        logger.info('Removing ambient RNA using CellSweep')
-        results['swept_counts'] = run_cellsweep(counts_dir=counts_dir, out_dir=out_dir, threads=threads, kwargs=kwargs)
 
     # Delete intermediate BUS files if requested
     if delete_bus:
@@ -1910,8 +1841,6 @@ def count_nac(
     quant_umis: bool = False,
     keep_flags: bool = False,
     exact_barcodes: bool = False,
-    remove_ambient: bool = False,
-    **kwargs
 ) -> Dict[str, Union[Dict[str, str], str]]:
     """Generates RNA velocity matrices for single-cell RNA seq.
 
@@ -1988,18 +1917,13 @@ def count_nac(
         quant_umis: Whether to run quant-tcc when there are UMIs, defaults to `False`
         keep_flags: Preserve flag column when sorting BUS file, defaults to `False`
         exact_barcodes: Use exact match for 'correcting' barcodes to on-list, defaults to `False`
-        remove_ambient: Whether to remove ambient RNA using CellSweep, defaults to `False`
-        **kwargs: Additional keyword arguments to pass to CellSweep
 
     Returns:
         Dictionary containing path to generated index
     """
     STATS.start()
     is_batch = isinstance(fastqs, str)
 
-    #* kwargs is only added for cellsweep, so check accordingly
-    check_kwargs(kwargs)
-
     results = {}
     make_directory(out_dir)
     unfiltered_results = results.setdefault('unfiltered', {})
@@ -2559,10 +2483,6 @@ def update_results_with_suffix(current_results, new_results, suffix):
             logger.warning(
                 'Plots for TCC matrices have not yet been implemented. The HTML report will not contain any plots.'
             )
-    
-    if remove_ambient:
-        logger.info('Removing ambient RNA using CellSweep')
-        results['swept_counts'] = run_cellsweep(counts_dir=counts_dir, out_dir=out_dir, threads=threads, kwargs=kwargs)
 
     # Delete intermediate BUS files if requested
     if delete_bus:
@@ -2621,8 +2541,6 @@ def count_velocity(
     strand: Optional[Literal['unstranded', 'forward', 'reverse']] = None,
     umi_gene: bool = False,
     em: bool = False,
-    remove_ambient: bool = False,
-    **kwargs
 ) -> Dict[str, Union[Dict[str, str], str]]:
     """Generates RNA velocity matrices (DEPRECATED).
 
@@ -2670,8 +2588,7 @@ def count_velocity(
             `False`
         em: Whether to estimate gene abundances using EM algorithm, defaults to
             `False`
-        remove_ambient: Whether to remove ambient RNA using CellSweep, defaults to `False`
-        **kwargs: Additional keyword arguments to pass to CellSweep
+
     Returns:
         Dictionary containing path to generated index
     """
@@ -2680,9 +2597,6 @@ def count_velocity(
     BUS_CDNA_PREFIX = 'spliced'
     BUS_INTRON_PREFIX = 'unspliced'
 
-    #* kwargs is only added for cellsweep, so check accordingly
-    check_kwargs(kwargs)
-
     results = {}
     make_directory(out_dir)
     unfiltered_results = results.setdefault('unfiltered', {})
@@ -2979,10 +2893,6 @@ def count_velocity(
     stats_path = STATS.save(os.path.join(out_dir, KB_INFO_FILENAME))
     results.update({'stats': stats_path})
 
-    if remove_ambient:
-        logger.info('Removing ambient RNA using CellSweep')
-        results['swept_counts'] = run_cellsweep(counts_dir=counts_dir, out_dir=out_dir, threads=threads, kwargs=kwargs)
-
     # Reports
     nb_path = os.path.join(out_dir, REPORT_NOTEBOOK_FILENAME)
     html_path = os.path.join(out_dir, REPORT_HTML_FILENAME)
@@ -3052,8 +2962,6 @@ def count_velocity_smartseq3(
     by_name: bool = False,
     inspect: bool = True,
     strand: Optional[Literal['unstranded', 'forward', 'reverse']] = None,
-    remove_ambient: bool = False,
-    **kwargs
 ) -> Dict[str, Union[str, Dict[str, str]]]:
     """Generates count matrices for Smartseq3 (DEPRECATED).
 
@@ -3080,8 +2988,6 @@ def count_velocity_smartseq3(
         inspect: Whether or not to inspect the output BUS file and generate
             the inspect.json
         strand: Strandedness, defaults to `None`
-        remove_ambient: Whether to remove ambient RNA using CellSweep, defaults to `False`
-        **kwargs: Additional keyword arguments to pass to CellSweep
 
     Returns:
         Dictionary containing paths to generated files
@@ -3091,9 +2997,6 @@ def count_velocity_smartseq3(
     BUS_CDNA_PREFIX = 'spliced'
     BUS_INTRON_PREFIX = 'unspliced'
 
-    #* kwargs is only added for cellsweep, so check accordingly
-    check_kwargs(kwargs)
-
     results = {}
     make_directory(out_dir)
     unfiltered_results = results.setdefault('unfiltered', {})