MadSpin: -O2 for productions with >= 4 decaying particles

MadSpin/decay.py

@@ -59,6 +59,7 @@
 #from madgraph.interface.madgraph_interface import MadGraphCmd
 import madgraph.interface.master_interface as Cmd
 import madgraph.interface.madevent_interface as me_interface
+import madgraph.interface.common_run_interface as common_run_interface
 import madgraph.iolibs.save_load_object as save_load_object
 import madgraph.iolibs.files as files
 import madgraph.fks.fks_common as fks_common
@@ -74,11 +75,40 @@
 
 MAX_COMPAT_FLAVS = 500
 
+# Optimisation flag the matrix elements are compiled with. It is the
+# GLOBAL_FLAG of the tree's Source/make_opts, which every makefile of the tree
+# (DHELAS, MODEL, the P directories, MadSpin/src/makefile_*) adds to FFLAGS;
+# ``output standalone`` leaves it empty, i.e. -O0.
+MS_GLOBAL_FLAG = '-O2'
+# ... used only when a production process has at least this many particles to
+# decay. -O2 costs ~2 s of compilation per run and only pays off when the
+# matrix elements dominate: measured on 10^4 events (the MG7 paper's MadSpin
+# timing setup), x1.05-1.07 for t t~ t t~ and x1.25 for W+W-W+W-, but 3-7%
+# slower for WH, t t~, ZZ and t t~ Z.
+MS_GLOBAL_FLAG_MIN_DECAYS = 4
+
 
 class MadSpinError(MadGraph5Error):
     pass
 
 
+def set_global_flag(tree, flag=MS_GLOBAL_FLAG):
+    """Have the standalone tree ``tree`` compiled with ``flag``.
+
+    ``output standalone`` has already built Source/DHELAS and Source/MODEL
+    with the empty default flag, and make does not track flags: when the flag
+    has to change, their objects and libraries are removed so that the next
+    make rebuilds them with it. Nothing is removed when it is already set.
+    """
+    make_opts = pjoin(tree, 'Source', 'make_opts')
+    if not common_run_interface.CommonRunCmd.update_make_opts_full(
+                                            make_opts, {'GLOBAL_FLAG': flag}):
+        return
+    for lib in ('DHELAS', 'MODEL'):
+        misc.compile(arg=['clean'], cwd=pjoin(tree, 'Source', lib),
+                     mode='fortran')
+
+
 def bw_retained_fraction(pole, width, bw_cut):
     """The fraction of a resonance's Breit-Wigner that a ``+- bw_cut * width``
     mass window keeps.
@@ -3959,6 +3989,24 @@ def choose_full_flavor_index(self, production_tag, decay, decay_me, event_map):
 
         return self.get_full_flavor_index(production_tag, decay_me, event_map)
 
+    def nb_decaying_in_production(self):
+        """The largest number of final-state particles of a production process
+        that MadSpin decays (the decays of their decay products do not count).
+        The production processes are in all_ME for madspin_v1 and in all_me
+        (type 'production') for the onshell/density modes."""
+        decay_ids = self.all_ME.decay_ids
+        finals = [topo['base_order'][1] for topo in self.all_ME.values()]
+        finals += [tag[1] for tag, info in getattr(self, 'all_me', {}).items()
+                   if info.get('type') == 'production']
+        return max([sum(abs(pid) in decay_ids for pid in final)
+                    for final in finals] or [0])
+
+    def use_global_flag(self, tree):
+        """Compile ``tree`` with MS_GLOBAL_FLAG if the production has enough
+        particles to decay for it to pay off (see MS_GLOBAL_FLAG_MIN_DECAYS)."""
+        if self.nb_decaying_in_production() >= MS_GLOBAL_FLAG_MIN_DECAYS:
+            set_global_flag(tree)
+
     def compile(self):
         logger.info('Compiling code')
         self.compile_fortran(self.path_me, mode="full_me")
@@ -3977,6 +4025,7 @@ def compile_fortran(self, path_me, mode='production_me'):
         logger.debug("""Finalizing %s's """% mode)
 
         # COMPILATION OF LIBRARY
+        self.use_global_flag(pjoin(path_me, mode))
         misc.compile( cwd=pjoin(path_me, mode,"Source","DHELAS"), mode='fortran')
         file_madspin=pjoin(MG5DIR, 'MadSpin', 'src', 'lha_read_ms.f')
         shutil.copyfile(file_madspin, pjoin(path_me, mode,"Source","MODEL","lha_read.f" ))
@@ -5644,13 +5693,15 @@ def compile(self):
             decay_args.insert(0, 'PROCNAME=_ms%d' % ms_run_id)
         #my_env = os.environ.copy()
         #os.environ["GFORTRAN_UNBUFFERED_ALL"] = "y"
+        self.use_global_flag(pjoin(self.path_me, ms_me_subdir))
         misc.compile(cwd=pjoin(self.path_me, ms_me_subdir, 'Source'),
                      nb_core=self.mgcmd.options['nb_core'])
         misc.compile(prod_args,
                      cwd=pjoin(self.path_me, ms_me_subdir, 'SubProcesses'),
                      nb_core=self.mgcmd.options['nb_core'])
         #Valentin: not sure the decay_folder exists in all cases, so I check
         if os.path.exists(pjoin(self.path_me, ms_me_decay_subdir)):
+            self.use_global_flag(pjoin(self.path_me, ms_me_decay_subdir))
             misc.compile(cwd=pjoin(self.path_me, ms_me_decay_subdir, 'Source'),
                         nb_core=self.mgcmd.options['nb_core'])
             misc.compile(decay_args,

madgraph/interface/common_run_interface.py

@@ -4778,6 +4778,7 @@ def update_make_opts_full(path, def_variables, keep_old=True):
         """update the make_opts file writing the environmental variables
         of def_variables.
         if a value of the dictionary is None then it is not written.
+        Return True if the file had to be changed.
         """
         make_opts = path
         pattern = re.compile(r'^(\w+)\s*=\s*(.*)$',re.DOTALL)
@@ -4868,7 +4869,7 @@ def update_make_opts_full(path, def_variables, keep_old=True):
             # never observe the file in the truncated state that open(...,'w')
             # would leave it in.
             misc.atomic_write(make_opts, content_variables + '\n'.join(content))
-        return       
+        return diff
 
 
 

tests/unit_tests/madspin/test_madspin.py

@@ -13268,3 +13268,97 @@ def test_without_a_deal_the_z_keys_are_the_occurrences(self):
 
 def _mom(part):
     return (part.E, part.px, part.py, part.pz)
+
+
+class TestSetGlobalFlag(unittest.TestCase):
+    """MadSpin compiles its standalone trees with madspin.MS_GLOBAL_FLAG; the
+    libraries ``output standalone`` already built must then be rebuilt."""
+
+    def setUp(self):
+        import tempfile
+        self.tree = tempfile.mkdtemp(prefix='madspin_flag_')
+        os.mkdir(pjoin(self.tree, 'Source'))
+        shutil.copy(pjoin(MG5DIR, 'Template', 'LO', 'Source', '.make_opts'),
+                    pjoin(self.tree, 'Source', 'make_opts'))
+        # stands in for the objects 'output standalone' left behind
+        for lib in ('DHELAS', 'MODEL'):
+            os.mkdir(pjoin(self.tree, 'Source', lib))
+            with open(pjoin(self.tree, 'Source', lib, 'makefile'), 'w') as fsock:
+                fsock.write('clean:\n\trm -f stale.o\n')
+
+    def tearDown(self):
+        shutil.rmtree(self.tree)
+
+    def _touch_objects(self):
+        for lib in ('DHELAS', 'MODEL'):
+            open(pjoin(self.tree, 'Source', lib, 'stale.o'), 'w').close()
+
+    def _objects(self):
+        return [os.path.exists(pjoin(self.tree, 'Source', lib, 'stale.o'))
+                for lib in ('DHELAS', 'MODEL')]
+
+    def _global_flag(self):
+        with open(pjoin(self.tree, 'Source', 'make_opts')) as fsock:
+            text = fsock.read()
+        head = text.split('#end_of_make_opts_variables')[0]
+        return [l.split('=', 1)[1] for l in head.splitlines()
+                if l.startswith('GLOBAL_FLAG=')]
+
+    def test_flag_is_written_and_stale_objects_removed(self):
+        self._touch_objects()
+        madspin.set_global_flag(self.tree)
+        self.assertEqual(self._global_flag(), [madspin.MS_GLOBAL_FLAG])
+        self.assertEqual(self._objects(), [False, False])
+
+    def test_nothing_is_removed_once_the_flag_is_set(self):
+        madspin.set_global_flag(self.tree)
+        self._touch_objects()
+        madspin.set_global_flag(self.tree)
+        self.assertEqual(self._global_flag(), [madspin.MS_GLOBAL_FLAG])
+        self.assertEqual(self._objects(), [True, True])
+
+    def test_a_different_flag_rebuilds_again(self):
+        madspin.set_global_flag(self.tree)
+        self._touch_objects()
+        madspin.set_global_flag(self.tree, '-O0')
+        self.assertEqual(self._global_flag(), ['-O0'])
+        self.assertEqual(self._objects(), [False, False])
+
+
+class TestGlobalFlagGate(unittest.TestCase):
+    """MS_GLOBAL_FLAG is only used when the production has at least
+    MS_GLOBAL_FLAG_MIN_DECAYS particles to decay."""
+
+    def _events(self, decay_ids):
+        obj = madspin.decay_all_events.__new__(madspin.decay_all_events)
+        obj.all_ME = madspin.AllMatrixElement(None, {}, decay_ids, None)
+        return obj
+
+    def test_count_from_madspin_v1_topologies(self):
+        obj = self._events([6, -6])
+        obj.all_ME[((21, 21), (6, -6, 6, -6))] = {'base_order': ((21, 21), (6, -6, 6, -6))}
+        obj.all_ME[((21, 21), (6, -6, 23))] = {'base_order': ((21, 21), (6, -6, 23))}
+        self.assertEqual(obj.nb_decaying_in_production(), 4)
+
+    def test_count_from_density_production_only(self):
+        # the decay processes (type 'decay') must not be counted, and the
+        # decays of decay products (the W of a top) are not production decays
+        obj = self._events([6, -6, 23])
+        obj.all_me = {((21, 21), (-6, 6, 23)): {'type': 'production'},
+                      ((6,), (5, 24)): {'type': 'decay'}}
+        self.assertEqual(obj.nb_decaying_in_production(), 3)
+
+    def test_flag_only_from_the_threshold_on(self):
+        calls = []
+        orig = madspin.set_global_flag
+        madspin.set_global_flag = lambda tree, *a, **k: calls.append(tree)
+        try:
+            for ndecay in range(1, madspin.MS_GLOBAL_FLAG_MIN_DECAYS + 2):
+                obj = self._events([24, -24])
+                obj.all_me = {((2, -1), (24,) * ndecay): {'type': 'production'}}
+                obj.use_global_flag('tree%d' % ndecay)
+        finally:
+            madspin.set_global_flag = orig
+        self.assertEqual(calls, ['tree%d' % n for n in
+                                 range(madspin.MS_GLOBAL_FLAG_MIN_DECAYS,
+                                       madspin.MS_GLOBAL_FLAG_MIN_DECAYS + 2)])