parallelised cspull

Luigi Tibaldo
1 parent 92a76a88
Showing 3 changed files with 121 additions and 5 deletions
cscripts/cspull.par
cscripts/cspull.py
test/test_cspull.py
@@ -50,8 +50,9 @@ seed,      i, h, 1,,, &quot;Initial random number generator seed&quot;
 #
 # Standard parameters
 #====================
-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, cspull.log,,, "Log filename"
+nthreads, i, h, 0,0,1000, "Number of parallel processes (0=use all available CPUs)"
+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, cspull.log,,, "Log filename"
@@ -23,6 +23,33 @@ import gammalib
 import ctools
 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, 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
+    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
  
  
 # ============ #
@@ -40,9 +67,43 @@ class cspull(ctools.csobservation):
         # Initialise application by calling the appropriate class constructor
         self._init_csobservation(self.__class__.__name__, ctools.__version__, argv)
  
+        # Initialise some members
+        self._nthreads = 0
+
         # 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),
+                 '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._nthreads = state['nthreads']
+
+        # Return
+        return
  
     # Private methods
     def _get_parameters(self):
@@ -92,6 +153,9 @@ class cspull(ctools.csobservation):
         #  Write input parameters into logger
         self._log_parameters(gammalib.TERSE)
  
+        # Set number of processes for multiprocessing
+        self._nthreads = mputils.nthreads(self)
+
         # Return
         return
  
@@ -310,6 +374,29 @@ 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)
+
+            # 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)
+
         # Loop over trials
         for seed in range(self['ntrials'].integer()):
  
@@ -260,6 +260,34 @@ class Test(test):
         # Check pull distribution file
         self._check_pull_file('cspull_py6.dat', rows=2)
  
+        # Set-up cspull without multiprocessing
+        pull = cscripts.cspull()
+        pull['inmodel']  = self._model
+        pull['onsrc']    = 'NONE'
+        pull['outfile']  = 'cspull_py7.dat'
+        pull['ntrials']  = 2
+        pull['caldb']    = self._caldb
+        pull['irf']      = self._irf
+        pull['ra']       = 83.6331
+        pull['dec']      = 22.0145
+        pull['emin']     = 0.1
+        pull['emax']     = 100.0
+        pull['enumbins'] = 0
+        pull['tmin']     = 0.0
+        pull['tmax']     = 100.0
+        pull['deadc']    = 0.98
+        pull['rad']      = 5.0
+        pull['logfile']  = 'cspull_py7.log'
+        pull['chatter']  = 2
+        pull['nthreads'] = 1
+
+        # Run cspull script
+        pull.logFileOpen()   # Make sure we get a log file
+        pull.run()
+
+        # Check pull distribution file
+        self._check_pull_file('cspull_py7.dat')
+
         # Return
         return