diff --git a/release-notes/next-release.md b/release-notes/next-release.md index 8bb9bbe10..0d21e7f12 100644 --- a/release-notes/next-release.md +++ b/release-notes/next-release.md @@ -3,6 +3,7 @@ ## New features * View number of genes in model notebook representation. +* Add progress bar to `flux_variability_analysis`. ## Fixes diff --git a/src/cobra/flux_analysis/variability.py b/src/cobra/flux_analysis/variability.py index 4b1066f2b..1fa2bea12 100644 --- a/src/cobra/flux_analysis/variability.py +++ b/src/cobra/flux_analysis/variability.py @@ -1,6 +1,5 @@ """Provide variability based methods such as flux variability or gene essentiality.""" - import logging from typing import TYPE_CHECKING, List, Optional, Set, Tuple, Union from warnings import warn @@ -8,6 +7,7 @@ import numpy as np import pandas as pd from optlang.symbolics import Zero +from rich.progress import Progress from ..core import Configuration, get_solution from ..util import ProcessPool @@ -26,6 +26,29 @@ configuration = Configuration() +class _FakeProgress: + def __enter__(self): + return self + + def __exit__(self, *args): + pass + + def add_task(self, *args, **kwargs): + return object() + + def update(*args, **kwargs): + pass + + +def _update_progress_bar( + progress_bar, + task, + advance: int, +) -> None: + """Update progress bar.""" + progress_bar.update(task, advance=advance) + + def _init_worker(model: "Model", loopless: bool, sense: str) -> None: """Initialize a global model object for multiprocessing. @@ -230,25 +253,43 @@ def flux_variability_analysis( model.add_cons_vars([flux_sum, flux_sum_constraint]) model.objective = Zero # This will trigger the reset as well - for what in ("minimum", "maximum"): - if processes > 1: - # We create and destroy a new pool here in order to set the - # objective direction for all reactions. This creates a - # slight overhead but seems the most clean. - chunk_size = len(reaction_ids) // processes - with ProcessPool( - processes, - initializer=_init_worker, - initargs=(model, loopless, what[:3]), - ) as pool: - for rxn_id, value in pool.imap_unordered( - _fva_step, reaction_ids, chunksize=chunk_size - ): + + _Progress = Progress + if logger.level != logging.INFO: + _Progress = _FakeProgress + + with _Progress() as progress: + progress_task = progress.add_task( + "[cyan]Performing FVA...", total=num_reactions + ) + + for what in ("minimum", "maximum"): + if processes > 1: + # We create and destroy a new pool here in order to set the + # objective direction for all reactions. This creates a + # slight overhead but seems the most clean. + chunk_size = len(reaction_ids) // processes + with ProcessPool( + processes, + initializer=_init_worker, + initargs=(model, loopless, what[:3]), + ) as pool: + for i, (rxn_id, value) in enumerate( + pool.imap_unordered( + _fva_step, reaction_ids, chunksize=chunk_size + ) + ): + fva_result.at[rxn_id, what] = value + _update_progress_bar( + progress, progress_task, (i + 1) / num_reactions + ) + else: + _init_worker(model, loopless, what[:3]) + for i, (rxn_id, value) in enumerate(map(_fva_step, reaction_ids)): fva_result.at[rxn_id, what] = value - else: - _init_worker(model, loopless, what[:3]) - for rxn_id, value in map(_fva_step, reaction_ids): - fva_result.at[rxn_id, what] = value + _update_progress_bar( + progress, progress_task, (i + 1) / num_reactions + ) return fva_result[["minimum", "maximum"]]