parallelized csspec

Luigi Tibaldo
1 parent b3f8e6c6
Showing 3 changed files with 136 additions and 14 deletions
cscripts/csspec.par
cscripts/csspec.py
test/test_csspec.py
@@ -46,9 +46,10 @@ fix_bkg,   b, h, no,,, &quot;Fix background parameters&quot;
 #
 # Standard parameters
 #====================
-publish, b, h, no,,, "Publish spectrum on VO Hub?"
-chatter, i, h, 2,0,4, "Chattiness of output"
-clobber, b, h, yes,,, "Overwrite existing output files with new output files?"
-debug,   b, h, no,,, "Activate debugging mode?"
-mode,    s, h, ql,,, "Mode of automatic parameters"
-logfile, f, h, csspec.log,,, "Log filename"
+nthreads, i, h, 0,0,1000, "Number of parallel processes (0=use all available CPUs)"
+publish,  b, h, no,,, "Publish spectrum on VO Hub?"
+chatter,  i, h, 2,0,4, "Chattiness of output"
+clobber,  b, h, yes,,, "Overwrite existing output files with new output files?"
+debug,    b, h, no,,, "Activate debugging mode?"
+mode,     s, h, ql,,, "Mode of automatic parameters"
+logfile,  f, h, csspec.log,,, "Log filename"
@@ -21,6 +21,33 @@
 import sys
 import gammalib
 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
+    buffer  = cls._log.buffer()
+
+    # Close logger
+    cls._log.close()
+
+    # Collect thread information
+    info = {'celapse': celapse, 'log': buffer}
+
+    # Return light curve bin result and thread information
+    return result, info
  
  
 # ============ #
@@ -48,6 +75,7 @@ class csspec(ctools.csobservation):
         self._binned_mode = False
         self._onoff_mode  = False
         self._method      = 'AUTO'
+        self._nthreads    = 0
  
         # Return
         return
@@ -59,6 +87,48 @@ class csspec(ctools.csobservation):
         # Return
         return
  
+    # State methods por pickling
+    def __getstate__(self):
+        """
+        Extend ctools.csobservation getstate method to include some members
+
+        Returns
+        -------
+        state : dict
+            Pickled instance
+        """
+        # Set pickled dictionary
+        state = {'base'         : ctools.csobservation.__getstate__(self),
+                 'ebounds'      : self._ebounds,
+                 'fits'         : self._fits,
+                 'binned_mode'  : self._binned_mode,
+                 'onoff_mode'   : self._onoff_mode,
+                 'method'       : self._method,
+                 'nthreads'     : self._nthreads}
+
+        # Return pickled dictionary
+        return state
+
+    def __setstate__(self, state):
+        """
+        Extend ctools.csobservation setstate method to include some members
+
+        Parameters
+        ----------
+        state : dict
+            Pickled instance
+        """
+        ctools.csobservation.__setstate__(self, state['base'])
+        self._ebounds       = state['ebounds']
+        self._fits          = state['fits']
+        self._binned_mode   = state['binned_mode']
+        self._onoff_mode    = state['onoff_mode']
+        self._method        = state['method']
+        self._nthreads      = state['nthreads']
+
+        # Return
+        return
+
  
     # Private methods
     def _get_parameters(self):
@@ -141,6 +211,9 @@ class csspec(ctools.csobservation):
         #  Write input parameters into logger
         self._log_parameters(gammalib.TERSE)
  
+        # Set number of processes for multiprocessing
+        self._nthreads = mputils.nthreads(self)
+
         # Write spectrum method header and parameters
         self._log_header1(gammalib.TERSE, 'Spectrum method')
         self._log_value(gammalib.TERSE, 'Unbinned CTA observations', n_unbinned)
@@ -575,6 +648,10 @@ class csspec(ctools.csobservation):
         result : dict
             Dictionary with fit results
         """
+
+        # Write header for energy bin
+        self._log_header2(gammalib.EXPLICIT, 'Energy bin ' + str(i + 1))
+
         # Get energy boundaries
         emin      = self._ebounds.emin(i)
         emax      = self._ebounds.emax(i)
@@ -692,17 +769,35 @@ class csspec(ctools.csobservation):
         # Initialise results
         results = []
  
-        # Loop over energy bins
-        for i in range(self._ebounds.size()):
+        # If using multiprocessing
+        if self._nthreads > 1:
  
-            # Write header for energy bin
-            self._log_header2(gammalib.EXPLICIT, 'Energy bin '+str(i+1))
+            # Create pool of workers
+            from multiprocessing import Pool
+            pool = Pool(processes = self._nthreads)
  
-            # Fit energy bin
-            result = self._fit_energy_bin(i)
+            # Run energy bin analysis in parallel with map
+            args        = [(self, i) for i in range(self._ebounds.size())]
+            poolresults = pool.map(_multiprocessing_func_wrapper, args)
  
-            # Append results
-            results.append(result)
+            # Close pool and join
+            pool.close()
+            pool.join()
+
+            # Construct results
+            for i in range(self._ebounds.size()):
+                results.append(poolresults[i][0])
+                self._log_string(gammalib.TERSE, poolresults[i][1]['log'], False)
+
+        # Otherwise, loop over energy bins
+        else:
+            for i in range(self._ebounds.size()):
+
+                # Fit energy bin
+                result = self._fit_energy_bin(i)
+
+                # Append results
+                results.append(result)
  
         # Return results
         return results
@@ -231,6 +231,32 @@ class Test(test):
         # Check result file
         self._check_result_file('csspec_py5.fits', 10)
  
+        # Set-up csspec without multiprocessing
+        spec = cscripts.csspec()
+        spec['inobs']    = self._events
+        spec['inmodel']  = self._model
+        spec['srcname']  = 'Crab'
+        spec['caldb']    = self._caldb
+        spec['irf']      = self._irf
+        spec['method']   = 'AUTO'
+        spec['ebinalg']  = 'LOG'
+        spec['enumbins'] = 5
+        spec['emin']     = 0.1
+        spec['emax']     = 100.0
+        spec['outfile']  = 'csspec_py6.fits'
+        spec['logfile']  = 'csspec_py6.log'
+        spec['chatter']  = 2
+        spec['publish']  = True
+        spec['nthreads'] = 1
+
+        # Run csspec script
+        spec.logFileOpen()   # Make sure we get a log file
+        spec.run()
+        spec.save()
+
+        # Check result file
+        self._check_result_file('csspec_py6.fits', 5)
+
         # Return
         return