exprmat 0.1.38__tar.gz → 0.1.40__tar.gz

{exprmat-0.1.38 → exprmat-0.1.40}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: exprmat
-Version: 0.1.38
+Version: 0.1.40
 Classifier: Programming Language :: Rust
 Classifier: Programming Language :: Python :: Implementation :: CPython
 Classifier: Programming Language :: Python :: Implementation :: PyPy
@@ -42,7 +42,6 @@ Requires-Dist: metacells
 Requires-Dist: sh
 Requires-Dist: networkx
 Requires-Dist: palantir
-Requires-Dist: rich
 Requires-Dist: adjusttext
 Requires-Dist: h5py
 Requires-Dist: cytoolz
@@ -55,10 +54,9 @@ Requires-Dist: hdbscan
 Requires-Dist: graphtools
 Requires-Dist: bincsync>=0.3.0
 Requires-Dist: pybigwig
+Requires-Dist: spatialdata
 Requires-Dist: loompy
-Requires-Dist: squidpy
 Requires-Dist: zarr
-Requires-Dist: spatialdata-io
 Requires-Dist: scikit-image
 Requires-Dist: spatial-image
 Requires-Dist: multiscale-spatial-image
@@ -74,11 +72,29 @@ Requires-Dist: formulaic>=1.0.2
 Requires-Dist: formulaic-contrasts>=0.2.0
 Requires-Dist: cython
 Requires-Dist: pyfastx
+Requires-Dist: genomicranges
+Requires-Dist: rustworkx
+Requires-Dist: xgboost
+Requires-Dist: dask-expr>=2.0
 Requires-Dist: sphinx ; extra == 'dev'
 Requires-Dist: twine ; extra == 'dev'
 Requires-Dist: ipykernel ; extra == 'dev'
 Requires-Dist: ipytree ; extra == 'dev'
+Requires-Dist: cudf-cu12==25.8.* ; extra == 'cuda12'
+Requires-Dist: dask-cudf-cu12==25.8.* ; extra == 'cuda12'
+Requires-Dist: cuml-cu12==25.8.* ; extra == 'cuda12'
+Requires-Dist: cugraph-cu12==25.8.* ; extra == 'cuda12'
+Requires-Dist: nx-cugraph-cu12==25.8.* ; extra == 'cuda12'
+Requires-Dist: cuxfilter-cu12==25.8.* ; extra == 'cuda12'
+Requires-Dist: cucim-cu12==25.8.* ; extra == 'cuda12'
+Requires-Dist: pylibraft-cu12==25.8.* ; extra == 'cuda12'
+Requires-Dist: raft-dask-cu12==25.8.* ; extra == 'cuda12'
+Requires-Dist: cuvs-cu12==25.8.* ; extra == 'cuda12'
+Requires-Dist: nx-cugraph-cu12==25.8.* ; extra == 'cuda12'
+Requires-Dist: cupy-cuda12x ; extra == 'cuda12'
+Requires-Dist: rapids-singlecell ; extra == 'cuda12'
 Provides-Extra: dev
+Provides-Extra: cuda12
 License-File: license
 Summary: exprmat: Routines for expression matrices
 Author-email: Zheng Yang <xornent@outlook.com>

{exprmat-0.1.38 → exprmat-0.1.40}/exprmat/__init__.py

@@ -1,10 +1,19 @@
 
+# the file specification, defined as the earlieast version to handle
+# current file standard. (revision number). specification versions are not
+# continuous. historically taken [1, 39]
+
+SPECIFICATION = 39
+
+import importlib.metadata
+ver_string = importlib.metadata.version("exprmat")
+MAJOR, MINOR, REVISION = [int(x) for x in ver_string.split('.')]
+
 import os
 os.environ['OPENBLAS_NUM_THREADS'] = '1'
 os.environ['GOTO_NUM_THREADS'] = '1'
 os.environ['OMP_NUM_THREADS'] = '1'
 
-import importlib.metadata
 import mudata as mu
 import anndata as ad
 import scanpy as sc
@@ -56,9 +65,11 @@ except ImportError:
     info('compilation finished.')
 
 
-from exprmat.configuration import default as config
+from exprmat.configuration import Configuration as configuration
 from exprmat.ansi import warning, info
 
+config = configuration()
+
 mu.set_options(pull_on_update = False)
 Image.MAX_IMAGE_PIXELS = config['max.image']
 
@@ -98,9 +109,11 @@ basepath = config['data']
 
 # core method exports
 from exprmat.utils import setup_styles
+from exprmat.ansi import error, warning, info, pprog, pproga
 from exprmat.reader.experiment import experiment, load_experiment
 from exprmat.reader.metadata import metadata, load_metadata
 
+
 setup_styles()
 
 
@@ -129,8 +142,8 @@ if PERFORM_DATABASE_CHECK:
 
 
 def cuda():
-
     import torch
+    from exprmat.utils import supports_tensorcore
     info(f'PyTorch version: {torch.__version__}')
     info(f'GPU acceleration availability: {"YES" if torch.cuda.is_available() else "NO"}')
     if torch.cuda.is_available():
@@ -144,7 +157,7 @@ def cuda():
             info(f'[{i}] {torch.cuda.get_device_name(i)} {"*" if torch.cuda.current_device() == i else ""}')
             info(f'    CUDA capability: {torch.cuda.get_device_capability(i)}')
             info(f'    Installed VRAM (GiB): {total_mem:.2f} GiB')
-            info(f'    Supporting TensorCore: {"YES" if torch.cuda.get_device_properties(0).major >= 7 else "NO"}')
+            info(f'    Supporting TensorCore: {"YES" if supports_tensorcore() else "NO"}')
             info(f'    Current dedicated memory: {dedicated_mem:.2f} / {total_mem:.2f} GiB ({100 * dedicated_mem / total_mem:.1f}%)')
 
 
@@ -154,10 +167,7 @@ def version():
     import platform
     from exprmat.ansi import info, error, format_file_size
 
-    ver_string = importlib.metadata.version("exprmat")
     db_ver = version_db()
-
-    MAJOR, MINOR, REVISION = [int(x) for x in ver_string.split('.')]
     info(f'exprmat {MAJOR}.{MINOR}.{REVISION} / exprmat-db {db_ver if db_ver is not None else "(Not installed)"}')
     info(f'os: {os.name} ({sys.platform})  platform version: {platform.release()}')
     info(f'loaded configuration from {finder}')
@@ -179,7 +189,6 @@ def memory():
 
 setwd = os.chdir
 getwd = os.getcwd
-def locate_data(dir): config.update_config('data', dir)
 
 
 __all__ = [
@@ -190,6 +199,12 @@ __all__ = [
     'load_metadata',
     'version',
     'memory',
+    'cuda',
     'setwd',
-    'getwd'
+    'getwd',
+    'error',
+    'warning',
+    'info',
+    'pprog',
+    'pproga'
 ]

{exprmat-0.1.38 → exprmat-0.1.40}/exprmat/ansi.py

@@ -2,6 +2,12 @@
 import re
 import os
 
+from functools import partial
+from tqdm import tqdm
+from tqdm.auto import tqdm as tqdma
+
+SILENT = False
+
 
 def fore_black() -> None:
     print('\033[30m', end = '')
@@ -107,6 +113,7 @@ def format_file_size(size):
 
 
 def error(text: str, error = None) -> None:
+    if SILENT: raise Exception(text) from error
     fore_red()
     print('[error]', end = ' ')
     ansi_reset()
@@ -117,6 +124,7 @@ def error(text: str, error = None) -> None:
 
 
 def warning(text: str) -> None:
+    if SILENT: return
     fore_yellow()
     print('[!]', end = ' ')
     ansi_reset()
@@ -125,6 +133,7 @@ def warning(text: str) -> None:
 
 
 def info(text: str) -> None:
+    if SILENT: return
     fore_cyan()
     print('[i]', end = ' ')
     ansi_reset()
@@ -147,3 +156,13 @@ def yellow(x): return '\033[33m' + x + '\033[0m'
 def blue(x): return '\033[34m' + x + '\033[0m'
 def purple(x): return '\033[35m' + x + '\033[0m'
 def cyan(x): return '\033[36m' + x + '\033[0m'
+
+
+progress_styles = {
+    'ncols': 80,
+    'ascii': '-━',
+    'bar_format': '   {bar} {desc:20} {n:5d} / {total:<5d} ({elapsed} < {remaining})'
+}
+
+pprog = partial(tqdm, **progress_styles)
+pproga = partial(tqdma, **progress_styles)

{exprmat-0.1.38 → exprmat-0.1.40}/exprmat/clustering/leiden.py

@@ -18,6 +18,8 @@ def leiden(
     neighbors_key: str | None = None,
     obsp: str | None = None,
     flavor = 'igraph',
+    use_gpu = False,
+    theta = 1,
     ** clustering_args,
 ) -> ad.AnnData | None:
     '''
@@ -86,11 +88,27 @@ def leiden(
         passes arguments to the `partition_type`) or `igraph.Graph.community_leiden` from `igraph`.
     '''
     
-    return sc.tl.leiden(
-        adata, resolution = resolution,
-        restrict_to = restrict_to, random_state = random_state,
-        key_added = key_added, adjacency = adjacency, directed = directed,
-        use_weights = use_weights, n_iterations = n_iterations,
-        partition_type = partition_type, neighbors_key = neighbors_key,
-        copy = False, flavor = flavor, **clustering_args
-    )
+    if use_gpu:
+        import rapids_singlecell as rsc
+        return rsc.tl.leiden(
+            adata, resolution = resolution,
+            restrict_to = restrict_to, random_state = random_state,
+            key_added = key_added, adjacency = adjacency,
+            use_weights = use_weights, 
+            # the n_iterations in igraph takes -1, 0, 1 and 2 as magic numbers
+            # not real number of iterations. indeed, the gpu implementation
+            # supports only real iterations, by default 200.
+            n_iterations = n_iterations if n_iterations > 2 else 200,
+            neighbors_key = neighbors_key, theta = theta,
+            copy = False
+        )
+
+    else:
+        return sc.tl.leiden(
+            adata, resolution = resolution,
+            restrict_to = restrict_to, random_state = random_state,
+            key_added = key_added, adjacency = adjacency, directed = directed,
+            use_weights = use_weights, n_iterations = n_iterations,
+            partition_type = partition_type, neighbors_key = neighbors_key,
+            copy = False, flavor = flavor, **clustering_args
+        )

{exprmat-0.1.38 → exprmat-0.1.40}/exprmat/cnv/infercnv.py

@@ -10,10 +10,9 @@ import scipy.ndimage
 import scipy.sparse
 from anndata import AnnData
 from scanpy import logging
-from tqdm.auto import tqdm
 from tqdm.contrib.concurrent import process_map
 
-from exprmat.ansi import error, warning, info
+from exprmat.ansi import error, warning, info, pproga
 from exprmat.utils import ensure_array
 
 
@@ -80,7 +79,7 @@ def infercnv(
     exclude_chromosomes
         List of chromosomes to exclude. The default is to exclude genosomes.
         For list of all available chromosomes in the dataset, you should manually check
-        the adata.var['.seqid'] column, when you are not using the pre-built genome annotations.
+        the adata.var['chr'] column, when you are not using the pre-built genome annotations.
 
     chunksize
         Process dataset in chunks of cells. This allows to run infercnv on
@@ -113,16 +112,16 @@ def infercnv(
 
     if not adata.var_names.is_unique:
         error('ensure your var_names are unique.')
-    if {'.seqid', '.start', '.end'} - set(adata.var.columns) != set():
-        warning('ensure your var table contains ".seqid", ".start", and ".end"')
+    if {'chr', 'start', 'end'} - set(adata.var.columns) != set():
+        warning('ensure your var table contains "chr", "start", and "end"')
         warning('this should be automatically queried for you if you specify a legal taxa name.')
         error('otherwise, you should parse the gff and configure these information yourself.')
 
-    var_mask = adata.var['.seqid'].isnull()
+    var_mask = adata.var['chr'].isnull()
     if np.sum(var_mask):
         warning(f'skipped {np.sum(var_mask)} genes because they don\'t have a genomic position annotated.')  # type: ignore
     if exclude_chromosomes is not None:
-        var_mask = var_mask | adata.var['.seqid'].isin(exclude_chromosomes)
+        var_mask = var_mask | adata.var['chr'].isin(exclude_chromosomes)
 
     valid_annot_adata = adata[:, ~ var_mask]
     reference = get_reference_cells(adata, reference_key, reference_cat, reference)[:, ~ var_mask]
@@ -131,7 +130,7 @@ def infercnv(
     if scipy.sparse.issparse(expr):
         expr = expr.tocsr()
 
-    var = valid_annot_adata.var.loc[:, ['.seqid', '.start', '.end']]
+    var = valid_annot_adata.var.loc[:, ['chr', 'start', 'end']]
     chr_pos, chunks, convolved_dfs = zip(
         *process_map(
             infercnv_proc,
@@ -143,7 +142,7 @@ def infercnv(
             itertools.repeat(step),
             itertools.repeat(dynamic_threshold),
             itertools.repeat(calculate_gene_values),
-            tqdm_class=tqdm,
+            tqdm_class = pproga,
             max_workers=cpu_count() if n_jobs is None else n_jobs,
         ),
         strict=False,
@@ -341,7 +340,7 @@ def running_mean_by_chr(
         and genomic position. This is a row in the plotted chromosome CNV matrix.
     '''
 
-    chromosomes = natural_sort([x for x in var['.seqid'].unique() if x != 'MT'])
+    chromosomes = natural_sort([x for x in var['chr'].unique() if x != 'MT'])
     running_means = [
         running_mean_chr(
             chr, expr, var, window_size, step, calculate_gene_values
@@ -361,7 +360,7 @@ def running_mean_by_chr(
 
 def running_mean_chr(chr, expr, var, window_size, step, calculate_gene_values):
 
-    genes = var.loc[var['.seqid'] == chr].sort_values('.start').index.values
+    genes = var.loc[var['chr'] == chr].sort_values('start').index.values
     tmp_x = expr[:, var.index.get_indexer(genes)]
     x_conv, convolved_gene_values = running_mean(
         tmp_x, n = window_size, step = step, 

{exprmat-0.1.38 → exprmat-0.1.40}/exprmat/configuration.py

@@ -17,25 +17,29 @@ class Configuration:
             'Ubuntu', 
             'Verdana'
         ]
-        
-        self.config['backend'] = 'TkAgg'
+
+        self.config['plotting.flavor'] = 'light'
+        self.config['plotting.backend'] = 'TkAgg'
         self.config['data'] = os.path.join(os.path.dirname(__file__), 'data')
 
         # target taxa
 
         self.config['taxa.reference'] = {
-            'mm10': 'mmu', # alias of grcm38
+            'mm10': 'mmu',   # the slightly modified version from ucsc based on grcm38
             'grcm39': 'mmu',
             'grcm38': 'mmu',
 
-            'hg19': 'hsa', # alias of grch37
+            'hg19': 'hsa',   # the slightly modified version from ucsc based on grch37
             'grch37': 'hsa',
             'grch38': 'hsa',
+
+            'dm6': 'dme'
         }
 
         self.config['default.assembly'] = {
             'mmu': 'grcm39',
-            'hsa': 'grch38'
+            'hsa': 'grch38',
+            'dme': 'dm6'
         }
 
         self.config['max.image'] = int(10000 * 50000)
@@ -53,6 +57,3 @@ class Configuration:
 
     def update_config(self, conf_name, value):
         self.config[conf_name] = value
-
-
-default = Configuration()

exprmat-0.1.40/exprmat/data/cre.py

@@ -0,0 +1,111 @@
+
+import os
+import pandas as pd
+import pickle
+import numpy as np
+import lightmotif
+
+from exprmat.data.finders import basepath
+from exprmat.ansi import error, warning, info
+
+with open(os.path.join(basepath, 'shared', 'motifs.pkl'), 'rb') as f:
+    motif_tree = pickle.load(f)
+
+pssm_table = {}
+transcription_factor_bindings = {}
+motif_subsets = {}
+
+
+def get_motifs(subset = 'all'):
+    if subset == 'all': return motif_tree
+    elif subset in motif_subsets.keys(): return motif_subsets[subset]
+    else:
+        subtree = {}
+        for key in motif_tree.keys():
+            items = {}
+            for mot in motif_tree[key].keys():
+                if motif_tree[key][mot]['source'] == subset:
+                    items[mot] = motif_tree[key][mot]
+            subtree[key] = items
+        motif_subsets[subset] = subtree
+        return motif_subsets[subset]
+
+
+def get_transcription_factor_binding_motifs(taxa):
+
+    if not taxa in transcription_factor_bindings.keys():
+        transcription_factor_bindings[taxa] = \
+            pd.read_feather(os.path.join(basepath, taxa, 'tf', 'transcription-factors.feather'))
+    
+    return transcription_factor_bindings[taxa]
+
+
+def search_transcription_factor(taxa, factor, subset = 'all'):
+
+    tfbindings = get_transcription_factor_binding_motifs(taxa)
+    if factor in tfbindings['symbol'].tolist():
+        row = tfbindings['symbol'].tolist().index(factor)
+        row = tfbindings.iloc[row, :].copy()
+        metaclust = row['metacluster'].split(' ')
+        motifs = row['motif'].split(' ')
+
+        table = get_motifs(subset)
+        pairs = []
+        for mc in metaclust:
+            for mot in motifs:
+                if mot in table[mc].keys():
+                    pairs.append((mc, mot))
+        
+        return pairs
+    
+    else: error(f'could not find transcription factor {factor} in taxa {taxa}.')
+
+
+def plot_sequence_logo(motif, fig_height = 1, fig_width_per_base = 0.2, dpi = 100):
+
+    from exprmat.plotting.sequence import sequence_logo
+    motifs = get_motifs()
+    if isinstance(motif, tuple):
+        return sequence_logo(
+            np.array(motifs[motif[0]][motif[1]]['ppm'].T, dtype = np.float32), 
+            ppm = True, y_format = 'bits', dpi = dpi, 
+            fig_height = fig_height, fig_width_per_base = fig_width_per_base
+        )
+    
+    elif isinstance(motif, str):
+        for mc in motifs.keys():
+            if motif in motifs[mc].keys():
+                return sequence_logo(
+                    np.array(motifs[mc][motif]['ppm'].T, dtype = np.float32), 
+                    ppm = True, y_format = 'bits', dpi = dpi, 
+                    fig_height = fig_height, fig_width_per_base = fig_width_per_base
+                )
+            
+
+def motif_to_tuple(motif):
+    motifs = get_motifs()
+    for mc in motifs.keys():
+        if motif in motifs[mc].keys():
+            return (mc, motif)
+    error(f'could not find motif {motif}')
+
+
+def to_pssm(motif, eps = 0.01):
+    
+    if isinstance(motif, str): motif = motif_to_tuple(motif)
+    if motif[1] in pssm_table.keys(): return pssm_table[motif[1]]
+
+    motifs = get_motifs()
+    integral = np.array(motifs[motif[0]][motif[1]]['ppm'].T * 100000, dtype = np.int32)
+    
+    # original order: acgt.
+    pwm = lightmotif.CountMatrix({
+        'A': integral[:, 0].tolist(),
+        'C': integral[:, 1].tolist(),
+        'T': integral[:, 3].tolist(),
+        'G': integral[:, 2].tolist(),
+    }).normalize(eps)
+
+    pssm = pwm.log_odds()
+    pssm_table[motif[1]] = pssm
+    return pssm

{exprmat-0.1.38 → exprmat-0.1.40}/exprmat/data/finders.py

@@ -3,7 +3,7 @@ import os
 import pandas
 from exprmat import basepath
 from exprmat.ansi import warning, info
-from exprmat.configuration import default as cfg
+from exprmat import config as cfg
 
 genome = {}
 
@@ -28,23 +28,28 @@ def get_genome(taxa):
 
     if taxa not in genome.keys(): genome[taxa] = {}
     if 'genes' in genome[taxa].keys(): return genome[taxa]['genes']
-    genome[taxa]['genes'] = pandas.read_table(
-        os.path.join(basepath, f'{taxa}', 'genome.tsv.gz'),
-        # set low memory to false to allow correct adjustment to mixed dtype.
-        index_col = '.ugene', dtype = {'.seqid': str}, low_memory = False
+    genome[taxa]['genes'] = pandas.read_feather(
+        os.path.join(basepath, f'{taxa}', 'genome.feather')
     )
 
+    genome[taxa]['genes'].index = genome[taxa]['genes']['uid'].tolist()
     return genome[taxa]['genes']
 
 
-def get_genome_size(assembly):
+def get_genome_size(assembly, as_dataframe = False):
 
     import json
     taxa = cfg['taxa.reference'][assembly.lower()]
     path = os.path.join(basepath, taxa, 'assemblies', assembly.lower(), 'size.json')
     with open(path, 'r', encoding = 'utf-8') as f:
         sizes = json.load(f)
-    return sizes
+    
+    if as_dataframe: return pandas.DataFrame({
+        'seqname': sizes.keys(),
+        'len': [sizes[k] for k in sizes.keys()]
+    })
+
+    else: return sizes
 
 
 def get_genome_gff_fname(assembly):
@@ -80,12 +85,13 @@ def get_genome_ranges(assembly, fname, as_grange = False):
     taxa = cfg['taxa.reference'][assembly.lower()]
     path = os.path.join(basepath, taxa, 'assemblies', assembly.lower(), f'{fname}.feather')
     feather = pandas.read_feather(path)
+    feather = feather.loc[feather['end'] - feather['start'] > 0, :].copy()
     
     if as_grange: 
         import genomicranges
         return genomicranges.GenomicRanges.from_pandas(
-            feather[['ucsc', 'start', 'end', 'strand', 'gene']].rename(
-            columns = {'ucsc': 'seqnames', 'start': 'starts', 'end': 'ends'}
+            feather[['chr', 'start', 'end', 'strand', 'gene']].rename(
+            columns = {'chr': 'seqnames', 'start': 'starts', 'end': 'ends'}
         ))
     else: return feather
 
@@ -185,7 +191,7 @@ def get_mapper_ensembl(taxa):
     else:
 
         gtable = get_genome(taxa)
-        ensembl_list = gtable[['ensembl']].values.transpose()[0].tolist()
+        ensembl_list = gtable[['id']].values.transpose()[0].tolist()
         id_list = gtable.index.tolist()
         ensembl_finder = {}
 

{exprmat-0.1.38 → exprmat-0.1.40}/exprmat/data/geneset.py

@@ -4,7 +4,7 @@ import pandas
 import numpy as np
 
 from exprmat.ansi import warning, info
-from exprmat.configuration import default as cfg
+from exprmat import config as cfg
 from exprmat.data.finders import basepath, genome, get_genome
 
 
@@ -59,13 +59,13 @@ def translate_id(taxa, genes, idfrom = 'entrez', idto = 'ugene', keep_nones = Tr
 
         # construct multi-platform convertion table.
         gtable = get_genome(taxa)
-        gtable = gtable[['ensembl', 'gene']].copy()
+        gtable = gtable[['id', 'gene']].copy()
         gtable['ugene'] = gtable.index.tolist()
-        gtable.index = gtable['ensembl'].tolist()
+        gtable.index = gtable['id'].tolist()
         gtable['uppercase'] = gtable['gene'].str.upper().tolist()
 
         # deduplicate
-        gtable = gtable.loc[~ gtable['ensembl'].duplicated(), :].copy()
+        gtable = gtable.loc[~ gtable['id'].duplicated(), :].copy()
         gtable['ugene'] = taxa + ':' + gtable['ugene']
 
         entrez = pandas.read_table(

{exprmat-0.1.38 → exprmat-0.1.40}/exprmat/descriptive/aucell.py

@@ -9,7 +9,7 @@ import numpy as np
 import pandas as pd
 from boltons.iterutils import chunked
 from exprmat.descriptive.recovery import enrich_cells
-from rich.progress import track
+from exprmat import pprog
 
 # to reduce the memory footprint of a ranking matrix we use unsigned 32bit 
 # integers which provides a range from 0 through 4,294,967,295. this should be 
@@ -96,7 +96,7 @@ def aucell4r(
                     module.noweights() if noweights else module,
                     auc_threshold = auc_threshold,
                 )
-                for module in track(signatures, description = 'auc')
+                for module in pprog(signatures, desc = 'auc')
             ]
         ).unstack("regulon")
         aucs.columns = aucs.columns.droplevel(0)

{exprmat-0.1.38 → exprmat-0.1.40}/exprmat/dynamics/parallel.py

@@ -3,11 +3,10 @@ import os
 from multiprocessing import Manager
 from threading import Thread
 from joblib import delayed, Parallel
-from rich.progress import TextColumn, BarColumn, TimeRemainingColumn, TimeElapsedColumn
 
 import numpy as np
 from scipy.sparse import issparse, spmatrix
-from exprmat.ansi import error, warning, info
+from exprmat.ansi import error, warning, info, pprog
 
 
 def get_n_jobs(n_jobs):
@@ -67,10 +66,10 @@ def parallelize(
         Whether to show a progress bar.
     """
 
-    if show_progress_bar: from rich.progress import Progress as progress
+    if show_progress_bar: progress = pprog
     else: progress = None
 
-    def update(pbar, task, queue, n_total):
+    def update(pbar, queue, n_total):
         
         n_finished = 0
         while n_finished < n_total:
@@ -83,32 +82,17 @@ def parallelize(
             assert res in (None, (1, None), 1)  # (None, 1) means only 1 job
             if res == (1, None):
                 n_finished += 1
-                if pbar is not None: pbar.advance(task, advance = 1)
+                if pbar is not None: pbar.update(1)
             elif res is None: n_finished += 1
-            elif pbar is not None: pbar.advance(task, advance = 1)
-
-        if pbar is not None:
-            pbar.stop()
+            elif pbar is not None: pbar.update(1)
 
 
     def wrapper(*args, **kwargs):
 
         if pass_queue and show_progress_bar:
-            pbar = None if progress is None else progress(
-                TextColumn("[progress.description]"),
-                BarColumn(),
-                TextColumn("[progress.percentage]{task.percentage:>4.1f}%"),
-                TimeRemainingColumn(),
-                TimeElapsedColumn()
-            )
-
-            pbartask = None
-            if pbar is not None: 
-                pbartask = pbar.add_task(description = "", total = col_len)
-                pbar.start()
-
+            pbar = None if progress is None else progress(total = col_len)
             queue = Manager().Queue()
-            thread = Thread(target = update, args = (pbar, pbartask, queue, len(collections)))
+            thread = Thread(target = update, args = (pbar, queue, len(collections)))
             thread.start()
 
         else: pbar, queue, thread = None, None, None