Add mputils.process() function and reduce code duplication

Jürgen Knödlseder
1 parent d3d48015
Showing 9 changed files with 108 additions and 113 deletions
cscripts/cslightcrv.py
cscripts/csphagen.py
cscripts/csphasecrv.py
cscripts/cspull.py
cscripts/cssens.py
cscripts/csspec.py
cscripts/cstsdist.py
cscripts/mputils.py
test/test_ctbutterfly.py
@@ -647,20 +647,13 @@ class cslightcrv(ctools.csobservation):
         # Write header
         self._log_header1(gammalib.TERSE, 'Generate lightcurve')
-        # If using multiprocessing
+        # If more than a single thread is requested then use multiprocessing
         if self._nthreads > 1:
-            # Create pool of workers
-            from multiprocessing import Pool
-            pool = Pool(processes = self._nthreads)
-
-            # Run time bin analysis in parallel with map
+            # Compute time bins
             args        = [(self, i) for i in range(self._tbins.size())]
-            poolresults = pool.map(_multiprocessing_func_wrapper, args)
-
-            # Close pool and join
-            pool.close()
-            pool.join()
+            poolresults = mputils.process(self._nthreads,
+                                          _multiprocessing_func_wrapper, args)
             # Construct results
             results = []
@@ -578,17 +578,10 @@ class csphagen(ctools.csobservation):
         # If there is more than one observation and we use multiprocessing
         if self._nthreads > 1 and self.obs().size() > 1:
-            # Create pool of workers
-            from multiprocessing import Pool
-            pool = Pool(processes = self._nthreads)
-
-            # Run time bin analysis in parallel with map
+            # Compute observations
             args        = [(self, i) for i in range(self.obs().size())]
-            poolresults = pool.map(_multiprocessing_func_wrapper, args)
-
-            # Close pool and join
-            pool.close()
-            pool.join()
+            poolresults = mputils.process(self._nthreads,
+                                          _multiprocessing_func_wrapper, args)
             # Construct results
             for i in range(self.obs().size()):
@@ -493,20 +493,13 @@ class csphasecrv(ctools.csobservation):
         # Write header
         self._log_header1(gammalib.TERSE, 'Generate phase curve')
-        # If using multiprocessing
+        # If more than a single thread is requested then use multiprocessing
         if self._nthreads > 1:
-            # Create pool of workers
-            from multiprocessing import Pool
-            pool = Pool(processes = self._nthreads)
-
-            # Run time bin analysis in parallel with map
+            # Compute phase bins
             args        = [(self, phbin) for phbin in self._phbins]
-            poolresults = pool.map(_multiprocessing_func_wrapper, args)
-
-            # Close pool and join
-            pool.close()
-            pool.join()
+            poolresults = mputils.process(self._nthreads,
+                                          _multiprocessing_func_wrapper, args)
             # Construct results
             for i in range(len(self._phbins)):
@@ -377,40 +377,34 @@ class cspull(ctools.csobservation):
         # Write header
         self._log_header1(gammalib.TERSE, 'Generate pull distribution')
-        # If using multiprocessing
-        if self._nthreads > 1:
-
-            # Create pool of workers
-            from multiprocessing import Pool
-            pool = Pool(processes = self._nthreads)
-
-            # Run time bin analysis in parallel with map
-            args        = [(self, seed + self['seed'].integer())
-                           for seed in range(self['ntrials'].integer())]
-            poolresults = pool.map(_multiprocessing_func_wrapper, args)
+        # Get number of trials
+        ntrials = self['ntrials'].integer()
-            # Close pool and join
-            pool.close()
-            pool.join()
+        # Get seed value
+        seed = self['seed'].integer()
-            # Write results
-            for seed in range(self['ntrials'].integer()):
-                result = poolresults[seed][0]
-                ioutils.write_csv_row(self['outfile'].filename().url(), seed,
-                                      result['colnames'], result['values'])
-                self._log_string(gammalib.TERSE, poolresults[seed][1]['log'], False)
+        # If more than a single thread is requested then use multiprocessing
+        if self._nthreads > 1:
+            args        = [(self, i + seed) for i in range(ntrials)]
+            poolresults = mputils.process(self._nthreads,
+                                          _multiprocessing_func_wrapper, args)
-        # Otherwise, loop over trials
-        else:
+        # Continue with regular processing
+        for i in range(ntrials):
-            for seed in range(self['ntrials'].integer()):
+            # If multiprocessing was used then recover results and put them
+            # into the log file
+            if self._nthreads > 1:
+                result = poolresults[i][0]
+                self._log_string(gammalib.TERSE, poolresults[i][1]['log'], False)
-                # Make a trial and add initial seed
-                result = self._trial(seed + self['seed'].integer())
+            # ... otherwise make a trial now
+            else:
+                result = self._trial(i + seed)
-                # Write out trial result
-                ioutils.write_csv_row(self['outfile'].filename().url(), seed,
-                                      result['colnames'], result['values'])
+            # Write out trial result
+            ioutils.write_csv_row(self['outfile'].filename().url(), i,
+                                  result['colnames'], result['values'])
         # Return
         return
@@ -715,17 +715,10 @@ class cssens(ctools.csobservation):
         # If using multiprocessing
         if self._nthreads > 1 and self._ebounds.size() > 1:
-            # Create pool of workers
-            from multiprocessing import Pool
-            pool = Pool(processes = self._nthreads)
-
-            # Run energy bin analysis in parallel with map
+            # Compute phase bins
             args        = [(self, ieng) for ieng in range(self._ebounds.size())]
-            poolresults = pool.map(_multiprocessing_func_wrapper, args)
-
-            # Close pool and join
-            pool.close()
-            pool.join()
+            poolresults = mputils.process(self._nthreads,
+                                          _multiprocessing_func_wrapper, args)
             # Construct results
             for ieng in range(self._ebounds.size()):
@@ -771,20 +771,13 @@ class csspec(ctools.csobservation):
         # Initialise results
         results = []
-        # If using multiprocessing
+        # If more than a single thread is requested then use multiprocessing
         if self._nthreads > 1:
-            # Create pool of workers
-            from multiprocessing import Pool
-            pool = Pool(processes = self._nthreads)
-
-            # Run energy bin analysis in parallel with map
+            # Compute energy bins
             args        = [(self, i) for i in range(self._ebounds.size())]
-            poolresults = pool.map(_multiprocessing_func_wrapper, args)
-
-            # Close pool and join
-            pool.close()
-            pool.join()
+            poolresults = mputils.process(self._nthreads,
+                                          _multiprocessing_func_wrapper, args)
             # Construct results
             for i in range(self._ebounds.size()):
@@ -329,38 +329,30 @@ class cstsdist(ctools.csobservation):
         # Write header
         self._log_header1(gammalib.TERSE, 'Generate TS distribution')
-        # If using multiprocessing
-        if self._nthreads > 1:
-
-            # Create pool of workers
-            from multiprocessing import Pool
-            pool = Pool(processes = self._nthreads)
-
-            # Run time bin analysis in parallel with map
-            args        = [(self, seed) for seed in range(self['ntrials'].integer())]
-            poolresults = pool.map(_multiprocessing_func_wrapper, args)
+        # Get number of trials
+        ntrials = self['ntrials'].integer()
-            # Close pool and join
-            pool.close()
-            pool.join()
-
-            # Write results
-            for seed in range(self['ntrials'].integer()):
-                result = poolresults[seed][0]
-                ioutils.write_csv_row(self['outfile'].filename().url(), seed,
-                                      result['colnames'], result['values'])
-                self._log_string(gammalib.TERSE, poolresults[seed][1]['log'], False)
+        # If more than a single thread is requested then use multiprocessing
+        if self._nthreads > 1:
+            args        = [(self, i) for i in range(ntrials)]
+            poolresults = mputils.process(self._nthreads,
+                                          _multiprocessing_func_wrapper, args)
-        # Otherwise, loop over trials
-        else:
+        # Continue with regular processing
+        for i in range(ntrials):
-            for seed in range(self['ntrials'].integer()):
+            # If multiprocessing was used then recover results and put them
+            # into the log file
+            if self._nthreads > 1:
+                result = poolresults[i][0]
+                self._log_string(gammalib.TERSE, poolresults[i][1]['log'], False)
-                # Make a trial
-                result = self._trial(seed)
+            # ... otherwise make a trial now
+            else:
+                result = self._trial(i)
-                # Write out trial result
-                ioutils.write_csv_row(self['outfile'].filename().url(), seed,
+            # Write out trial result
+            ioutils.write_csv_row(self['outfile'].filename().url(), i,
                                   result['colnames'], result['values'])
         # Return
@@ -20,12 +20,12 @@
 import gammalib
-# ==================== #
-# Number of processes  #
-# ==================== #
+# ====================== #
+# Get number of threads  #
+# ====================== #
 def nthreads(cls):
     """
-    Determines the number of parallel processes to use.
+    Determines the number of parallel processes to use
     The number is based on the user parameter "nthreads" and the availability
     of the multiprocessing module.
@@ -38,7 +38,7 @@ def nthreads(cls):
     Returns
     -------
     nthreads : int
-        Number of parallel processes
+        Number of threads
     """
     # Log multiprocessing configuration
     cls._log_header1(gammalib.NORMAL, 'Multiprocessing')
@@ -69,3 +69,41 @@ def nthreads(cls):
     # Return number of processes to be used
     return nthreads
+
+
+# ============================ #
+# Execute function in parallel #
+# ============================ #
+def process(nthreads, function, args):
+    """
+    Execute function in parallel
+
+    Parameters
+    ----------
+    nthreads : int
+        Number of parallel threads
+    function : Python function
+        Function to be executed in parallel
+    args : list of tuples
+        Function arguments for each evaluation
+
+    Returns
+    -------
+    results : list of dicts
+        List of function evaluation results
+    """
+    # Import Pool module from multiprocessing module
+    from multiprocessing import Pool
+
+    # Create a pool of processes
+    pool = Pool(processes=nthreads)
+
+    # Run function in parallel
+    results = pool.map(function, args)
+
+    # Close pool and join
+    pool.close()
+    pool.join()
+
+    # Return results
+    return results
@@ -123,6 +123,12 @@ class Test(test):
         # Check result file
         self._check_result_file('ctbutterfly_py1.dat')
+        # Check CSV result
+        self.test_value(butterfly.butterfly().nrows(), 100,
+             'Check for 100 rows in CSV file')
+        self.test_value(butterfly.butterfly().ncols(), 4,
+             'Check for 4 columns in CSV file')
+
         # Set-up ctbutterfly
         butterfly = ctools.ctbutterfly()
         butterfly['inobs']   = self._events