From d5c0427a51c5a66219ed41a211215605d3180204 Mon Sep 17 00:00:00 2001
From: Giordon Stark <kratsg@gmail.com>
Date: Thu, 29 Oct 2020 16:52:50 -0400
Subject: [PATCH] allow for tqdm to work both ways

---
 src/pyhf/infer/calculators.py | 78 +++++++++++++++++++++++++----------
 1 file changed, 56 insertions(+), 22 deletions(-)

diff --git a/src/pyhf/infer/calculators.py b/src/pyhf/infer/calculators.py
index 6f03d15769..050eb8e497 100644
--- a/src/pyhf/infer/calculators.py
+++ b/src/pyhf/infer/calculators.py
@@ -540,11 +540,55 @@ def distributions(self, poi_test, track_progress=None):
             unit='toy',
         )
 
-        signal_teststat = self.executor.map(
-            teststat_func,
-            *zip(
-                *(
-                    (
+        if not isinstance(self.executor, futures.TrivialExecutor):
+            signal_teststat = self.executor.map(
+                teststat_func,
+                *zip(
+                    *(
+                        (
+                            poi_test,
+                            sample,
+                            self.pdf,
+                            signal_pars,
+                            self.par_bounds,
+                            self.fixed_params,
+                        )
+                        for sample in signal_sample
+                    )
+                ),
+            )
+
+            bkg_teststat = self.executor.map(
+                teststat_func,
+                *zip(
+                    *(
+                        (
+                            poi_test,
+                            sample,
+                            self.pdf,
+                            bkg_pars,
+                            self.par_bounds,
+                            self.fixed_params,
+                        )
+                        for sample in bkg_sample
+                    )
+                ),
+            )
+
+            signal_teststat = tqdm.tqdm(
+                signal_teststat, **tqdm_options, desc='Signal-like'
+            )
+            s_plus_b = EmpiricalDistribution(tensorlib.astensor(list(signal_teststat)))
+            bkg_teststat = tqdm.tqdm(
+                bkg_teststat, **tqdm_options, desc='Background-like'
+            )
+            b_only = EmpiricalDistribution(tensorlib.astensor(list(bkg_teststat)))
+        else:
+            signal_teststat = []
+            for sample in tqdm.tqdm(signal_sample, **tqdm_options, desc='Signal-like'):
+                signal_teststat.append(
+                    self.executor.submit(
+                        teststat_func,
                         poi_test,
                         sample,
                         self.pdf,
@@ -552,16 +596,13 @@ def distributions(self, poi_test, track_progress=None):
                         self.par_bounds,
                         self.fixed_params,
                     )
-                    for sample in signal_sample
                 )
-            ),
-        )
 
-        bkg_teststat = self.executor.map(
-            teststat_func,
-            *zip(
-                *(
-                    (
+            bkg_teststat = []
+            for sample in tqdm.tqdm(bkg_sample, **tqdm_options, desc='Background-like'):
+                bkg_teststat.append(
+                    self.executor.submit(
+                        teststat_func,
                         poi_test,
                         sample,
                         self.pdf,
@@ -569,18 +610,11 @@ def distributions(self, poi_test, track_progress=None):
                         self.par_bounds,
                         self.fixed_params,
                     )
-                    for sample in bkg_sample
                 )
-            ),
-        )
-
-        signal_teststat = tqdm.tqdm(signal_teststat, **tqdm_options, desc='Signal-like')
-
-        s_plus_b = EmpiricalDistribution(tensorlib.astensor(list(signal_teststat)))
 
-        bkg_teststat = tqdm.tqdm(bkg_teststat, **tqdm_options, desc='Background-like')
+            s_plus_b = EmpiricalDistribution(tensorlib.astensor(list(signal_teststat)))
+            b_only = EmpiricalDistribution(tensorlib.astensor(list(bkg_teststat)))
 
-        b_only = EmpiricalDistribution(tensorlib.astensor(list(bkg_teststat)))
         return s_plus_b, b_only
 
     def teststatistic(self, poi_test):