Reduce code duplication for multiprocessing in cscripts

Jürgen Knödlseder
1 parent 5aeb1560
Showing 8 changed files with 64 additions and 212 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
@@ -26,33 +26,6 @@ from cscripts import ioutils
 from cscripts import mputils
-# ============================================ #
-# Global functions for multiprocessing support #
-# ============================================ #
-def _multiprocessing_func_wrapper(args):
-    return _multiprocessing_func(*args)
-def _multiprocessing_func(cls, i):
-
-    # Initialise thread logger
-    cls._log.clear()
-    cls._log.buffer_size(100000)
-
-    # Compute light curve bin
-    cstart  = cls.celapse()
-    result  = cls._timebin(i)
-    celapse = cls.celapse() - cstart
-    log     = cls._log.buffer()
-
-    # Close logger
-    cls._log.close()
-
-    # Collect thread information
-    info = {'celapse': celapse, 'log': log}
-
-    # Return light curve bin result and thread information
-    return result, info
-
-
 # ================ #
 # cslightcrv class #
 # ================ #
@@ -651,9 +624,9 @@ class cslightcrv(ctools.csobservation):
         if self._nthreads > 1:
             # Compute time bins
-            args        = [(self, i) for i in range(self._tbins.size())]
-            poolresults = mputils.process(self._nthreads,
-                                          _multiprocessing_func_wrapper, args)
+            args        = [(self, '_timebin', i)
+                           for i in range(self._tbins.size())]
+            poolresults = mputils.process(self._nthreads, mputils.mpfunc, args)
             # Construct results
             results = []
@@ -25,32 +25,6 @@ import math
 import sys
 from cscripts import mputils
-# ============================================ #
-# Global functions for multiprocessing support #
-# ============================================ #
-def _multiprocessing_func_wrapper(args):
-    return _multiprocessing_func(*args)
-def _multiprocessing_func(cls, i):
-
-    # Initialise thread logger
-    cls._log.clear()
-    cls._log.buffer_size(100000)
-
-    # Compute light curve bin
-    cstart  = cls.celapse()
-    result  = cls._process_observation(i)
-    celapse = cls.celapse() - cstart
-    log     = cls._log.buffer()
-
-    # Close logger
-    cls._log.close()
-
-    # Collect thread information
-    info = {'celapse': celapse, 'log': log}
-
-    # Return light curve bin result and thread information
-    return result, info
-
 # =============== #
 # csfindobs class #
@@ -579,9 +553,9 @@ class csphagen(ctools.csobservation):
         if self._nthreads > 1 and self.obs().size() > 1:
             # Compute observations
-            args        = [(self, i) for i in range(self.obs().size())]
-            poolresults = mputils.process(self._nthreads,
-                                          _multiprocessing_func_wrapper, args)
+            args        = [(self, '_process_observation', i)
+                           for i in range(self.obs().size())]
+            poolresults = mputils.process(self._nthreads, mputils.mpfunc, args)
             # Construct results
             for i in range(self.obs().size()):
@@ -25,32 +25,6 @@ from cscripts import obsutils
 from cscripts import ioutils
 from cscripts import mputils
-# ============================================ #
-# Global functions for multiprocessing support #
-# ============================================ #
-def _multiprocessing_func_wrapper(args):
-    return _multiprocessing_func(*args)
-def _multiprocessing_func(cls, phbin):
-
-    # Initialise thread logger
-    cls._log.clear()
-    cls._log.buffer_size(100000)
-
-    # Compute light curve bin
-    cstart  = cls.celapse()
-    result  = cls._phase_bin(phbin)
-    celapse = cls.celapse() - cstart
-    log     = cls._log.buffer()
-
-    # Close logger
-    cls._log.close()
-
-    # Collect thread information
-    info = {'celapse': celapse, 'log': log}
-
-    # Return light curve bin result and thread information
-    return result, info
-
 # ================ #
 # csphasecrv class #
@@ -497,9 +471,8 @@ class csphasecrv(ctools.csobservation):
         if self._nthreads > 1:
             # Compute phase bins
-            args        = [(self, phbin) for phbin in self._phbins]
-            poolresults = mputils.process(self._nthreads,
-                                          _multiprocessing_func_wrapper, args)
+            args        = [(self, '_phase_bin', phbin) for phbin in self._phbins]
+            poolresults = mputils.process(self._nthreads, mputils.mpfunc, args)
             # Construct results
             for i in range(len(self._phbins)):
@@ -26,33 +26,6 @@ from cscripts import ioutils
 from cscripts import mputils
-# ============================================ #
-# Global functions for multiprocessing support #
-# ============================================ #
-def _multiprocessing_func_wrapper(args):
-    return _multiprocessing_func(*args)
-def _multiprocessing_func(cls, i):
-
-    # Initialise thread logger
-    cls._log.clear()
-    cls._log.buffer_size(100000)
-
-    # Compute light curve bin
-    cstart  = cls.celapse()
-    result  = cls._trial(i)
-    celapse = cls.celapse() - cstart
-    log     = cls._log.buffer()
-
-    # Close logger
-    cls._log.close()
-
-    # Collect thread information
-    info = {'celapse': celapse, 'log': log}
-
-    # Return light curve bin result and thread information
-    return result, info
-
-
 # ============ #
 # cspull class #
 # ============ #
@@ -385,9 +358,8 @@ class cspull(ctools.csobservation):
         # 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)
+            args        = [(self, '_trial', i + seed) for i in range(ntrials)]
+            poolresults = mputils.process(self._nthreads, mputils.mpfunc, args)
         # Continue with regular processing
         for i in range(ntrials):
@@ -28,33 +28,6 @@ from cscripts import ioutils
 from cscripts import mputils
-# ============================================ #
-# Global functions for multiprocessing support #
-# ============================================ #
-def _multiprocessing_func_wrapper(args):
-    return _multiprocessing_func(*args)
-def _multiprocessing_func(cls, i):
-
-    # Initialise thread logger
-    cls._log.clear()
-    cls._log.buffer_size(100000)
-
-    # Compute light curve bin
-    cstart  = cls.celapse()
-    result  = cls._e_bin(i)
-    celapse = cls.celapse() - cstart
-    log     = cls._log.buffer()
-
-    # Close logger
-    cls._log.close()
-
-    # Collect thread information
-    info = {'celapse': celapse, 'log': log}
-
-    # Return light curve bin result and thread information
-    return result, info
-
-
 # ============ #
 # cssens class #
 # ============ #
@@ -715,10 +688,10 @@ class cssens(ctools.csobservation):
         # If using multiprocessing
         if self._nthreads > 1 and self._ebounds.size() > 1:
-            # Compute phase bins
-            args        = [(self, ieng) for ieng in range(self._ebounds.size())]
-            poolresults = mputils.process(self._nthreads,
-                                          _multiprocessing_func_wrapper, args)
+            # Compute energy bins
+            args        = [(self, '_e_bin', i)
+                           for i in range(self._ebounds.size())]
+            poolresults = mputils.process(self._nthreads, mputils.mpfunc, args)
             # Construct results
             for ieng in range(self._ebounds.size()):
@@ -24,33 +24,6 @@ import ctools
 from cscripts import mputils
-# ============================================ #
-# Global functions for multiprocessing support #
-# ============================================ #
-def _multiprocessing_func_wrapper(args):
-    return _multiprocessing_func(*args)
-def _multiprocessing_func(cls, i):
-
-    # Initialise thread logger
-    cls._log.clear()
-    cls._log.buffer_size(100000)
-
-    # Compute light curve bin
-    cstart  = cls.celapse()
-    result  = cls._fit_energy_bin(i)
-    celapse = cls.celapse() - cstart
-    log     = cls._log.buffer()
-
-    # Close logger
-    cls._log.close()
-
-    # Collect thread information
-    info = {'celapse': celapse, 'log': log}
-
-    # Return light curve bin result and thread information
-    return result, info
-
-
 # ============ #
 # csspec class #
 # ============ #
@@ -775,9 +748,9 @@ class csspec(ctools.csobservation):
         if self._nthreads > 1:
             # Compute energy bins
-            args        = [(self, i) for i in range(self._ebounds.size())]
-            poolresults = mputils.process(self._nthreads,
-                                          _multiprocessing_func_wrapper, args)
+            args        = [(self, '_fit_energy_bin', i)
+                           for i in range(self._ebounds.size())]
+            poolresults = mputils.process(self._nthreads, mputils.mpfunc, args)
             # Construct results
             for i in range(self._ebounds.size()):
@@ -27,33 +27,6 @@ from cscripts import ioutils
 from cscripts import mputils
-# ============================================ #
-# Global functions for multiprocessing support #
-# ============================================ #
-def _multiprocessing_func_wrapper(args):
-    return _multiprocessing_func(*args)
-def _multiprocessing_func(cls, i):
-
-    # Initialise thread logger
-    cls._log.clear()
-    cls._log.buffer_size(100000)
-
-    # Compute light curve bin
-    cstart  = cls.celapse()
-    result  = cls._trial(i)
-    celapse = cls.celapse() - cstart
-    log     = cls._log.buffer()
-
-    # Close logger
-    cls._log.close()
-
-    # Collect thread information
-    info = {'celapse': celapse, 'log': log}
-
-    # Return light curve bin result and thread information
-    return result, info
-
-
 # ============== #
 # cstsdist class #
 # ============== #
@@ -334,9 +307,8 @@ class cstsdist(ctools.csobservation):
         # 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)
+            args        = [(self, '_trial', i) for i in range(ntrials)]
+            poolresults = mputils.process(self._nthreads, mputils.mpfunc, args)
         # Continue with regular processing
         for i in range(ntrials):
@@ -20,9 +20,9 @@
 import gammalib
-# ====================== #
-# Get number of threads  #
-# ====================== #
+# ===================== #
+# Get number of threads #
+# ===================== #
 def nthreads(cls):
     """
     Determines the number of parallel processes to use
@@ -107,3 +107,45 @@ def process(nthreads, function, args):
     # Return results
     return results
+
+
+# ======== #
+# Function #
+# ======== #
+def mpfunc(args):
+    """
+    Multiprocessing function
+
+    Parameters
+    ----------
+    args : tuple
+        Tuple comprised of class, function name and function argument
+
+    Returns
+    -------
+    result, info : tuple
+        Tuple of function result and information
+    """
+    # Extract
+    cls = args[0] # Class
+    fct = args[1] # Class method name
+    i   = args[2] # Function argument
+
+    # Initialise thread logger
+    cls._log.clear()
+    cls._log.buffer_size(100000)
+
+    # Compute function for argument
+    cstart  = cls.celapse()
+    result  = getattr(cls, fct)(i)
+    celapse = cls.celapse() - cstart
+    log     = cls._log.buffer()
+
+    # Close logger
+    cls._log.close()
+
+    # Collect thread information
+    info = {'celapse': celapse, 'log': log}
+
+    # Return light curve bin result and thread information
+    return result, info