single thread hybrid solver

YosefLab · Nov 16, 2023 · 0de7452 · 0de7452
1 parent 2ca07df
commit 0de7452
Showing 1 changed file with 33 additions and 20 deletions.
diff --git a/cassiopeia/solver/HybridSolver.py b/cassiopeia/solver/HybridSolver.py
@@ -144,27 +144,40 @@ 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]}-"
-                                        f"{next(logfile_names)}.log",
-                                layer,
-                            )
-                            for subproblem in subproblems
-                        ],
-                    ),
-                    total=len(subproblems),
+        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],
+                                    f"{logfile.split('.log')[0]}-"
+                                            f"{next(logfile_names)}.log",
+                                    layer,
+                                )
+                                for subproblem in subproblems
+                            ],
+                        ),
+                        total=len(subproblems),
+                    )
                 )
-            )
+        # single-threaded bottom solver approach
+        else:
+            results = [
+                self.apply_bottom_solver(
+                    cassiopeia_tree,
+                    subproblem[0],
+                    subproblem[1],
+                    f"{logfile.split('.log')[0]}-{next(logfile_names)}.log",
+                    layer,
+                )
+                for subproblem in tqdm(subproblems, total=len(subproblems))
+            ]
 
         for result in results: