Merge pull request #233 from YosefLab/single-thread-hybrid

Single thread hybrid solver
YosefLab · Feb 6, 2024 · 97accd7 · 97accd7
2 parents 5b8a452 + 42cd03f
commit 97accd7
Show file tree

Hide file tree

Showing 2 changed files with 55 additions and 25 deletions.
diff --git a/cassiopeia/solver/HybridSolver.py b/cassiopeia/solver/HybridSolver.py
@@ -59,6 +59,8 @@ class HybridSolver(CassiopeiaSolver.CassiopeiaSolver):
                 "inverse": Transforms each probability p by taking 1/p
                 "square_root_inverse": Transforms each probability by the
                     the square root of 1/p
+        progress_bar: Indicates if a progress bar should be shown when
+            `solve` is called.
     """
 
     def __init__(
@@ -69,6 +71,7 @@ def __init__(
         cell_cutoff: int = None,
         threads: int = 1,
         prior_transformation: str = "negative_log",
+        progress_bar: bool = True,
     ):
 
         if lca_cutoff is None and cell_cutoff is None:
@@ -89,6 +92,7 @@ def __init__(
         self.cell_cutoff = cell_cutoff
 
         self.threads = threads
+        self.progress_bar = progress_bar
 
     def solve(
         self,
@@ -144,27 +148,45 @@ def solve(
         logfile_names = iter([i for i in range(1, len(subproblems) + 1)])
 
         # multi-threaded bottom solver approach
-        with multiprocessing.Pool(processes=self.threads) as pool:
-
-            results = list(
-                tqdm(
-                    pool.starmap(
-                        self.apply_bottom_solver,
-                        [
-                            (
-                                cassiopeia_tree,
-                                subproblem[0],
-                                subproblem[1],
-                                f"{logfile.split('.log')[0]}-"
+        if self.threads > 1:
+            with multiprocessing.Pool(processes=self.threads) as pool:
+
+                results = list(
+                    tqdm(
+                        pool.starmap(
+                            self.apply_bottom_solver,
+                            [
+                                (
+                                    cassiopeia_tree,
+                                    subproblem[0],
+                                    subproblem[1],
+                                    None if logfile is None else 
+                                        f"{logfile.split('.log')[0]}-"
                                         f"{next(logfile_names)}.log",
-                                layer,
-                            )
-                            for subproblem in subproblems
-                        ],
-                    ),
-                    total=len(subproblems),
+                                    layer,
+                                )
+                                for subproblem in subproblems
+                            ],
+                        ),
+                        total=len(subproblems),
+                        disable=not self.progress_bar
+                    )
                 )
-            )
+        # single-threaded bottom solver approach
+        else:
+            results = [
+                self.apply_bottom_solver(
+                    cassiopeia_tree,
+                    subproblem[0],
+                    subproblem[1],
+                    None if logfile is None else 
+                        f"{logfile.split('.log')[0]}-"
+                        f"{next(logfile_names)}.log",
+                    layer,
+                )
+                for subproblem in tqdm(subproblems, 
+                    total=len(subproblems),disable=not self.progress_bar)
+            ]
 
         for result in results:
 

diff --git a/cassiopeia/solver/ILPSolver.py b/cassiopeia/solver/ILPSolver.py
@@ -126,10 +126,14 @@ def solve(
                 " analysis."
             )
 
-        # setup logfile config
-        handler = logging.FileHandler(logfile)
-        handler.setLevel(logging.INFO)
-        logger.addHandler(handler)
+        # configure logger
+        if logfile is not None:
+            file_handler = logging.FileHandler(logfile)
+            file_handler.setLevel(logging.INFO)
+            logger.addHandler(file_handler)
+        logger.ch.setLevel(logging.getLogger().level)
+
+        # add to logger
         logger.info("Solving tree with the following parameters.")
         logger.info(f"Convergence time limit: {self.convergence_time_limit}")
         logger.info(
@@ -272,7 +276,10 @@ def solve(
             cassiopeia_tree.collapse_mutationless_edges(
                 infer_ancestral_characters=True
             )
-        logger.removeHandler(handler)
+
+        if logfile is not None:
+            logger.removeHandler(file_handler)
+            file_handler.close()
 
     def infer_potential_graph(
         self,
@@ -534,7 +541,8 @@ def solve_steiner_instance(
 
         # Add user-defined parameters
         model.params.MIPGAP = self.mip_gap
-        model.params.LogFile = logfile
+        if logfile is not None:
+            model.params.LogFile = logfile
 
         if self.seed is not None:
             model.params.Seed = self.seed