MadSpin: six fixes for the frame-projected polarised density

MadSpin/interface_madspin.py

@@ -8591,8 +8591,16 @@ def _production_polarization(self):
             source = {}
             multi_id = set()
             for line in lines:
+                # The perturbation-order bracket of an NLO process line
+                # ('p p > z{0} z{0} [QCD]') is irrelevant to the polarisation
+                # braces, which sit on the legs -- but MadSpin's mg5cmd holds
+                # the TREE-level model, so extract_process refuses it with
+                # "Perturbation order QCD is not among the perturbation orders
+                # allowed for by the loop model" and every NLO polarised
+                # production silently lost its restriction.  Strip it.
+                bare = re.sub(r'\[[^\]]*\]', '', line).strip()
                 try:
-                    procdef = self.mg5cmd.extract_process(line)
+                    procdef = self.mg5cmd.extract_process(bare)
                 except Exception as error:
                     logger.warning('MadSpin could not re-read the polarisation of '
                                    'the production process "%s" (%s); the density '
@@ -9916,6 +9924,10 @@ def _onshell_production_norm(self, production, prod_static):
         Evaluated on a round-tripped copy, like ``_upfront_production``'s
         ``prod_off``, so the two sides of the ratio see the same %.10e
         truncation.
+
+        Any residual inaccuracy of the boosted density evaluation is already
+        inside the numerator of every ``me_frame`` run; taking the denominator
+        in the same frame is what makes it cancel.
         """
         frame_boost = self._frame_boost(production)
         if frame_boost is None:
@@ -10630,11 +10642,13 @@ def sequential_accept_reject(self, production, evt_decayfile, maxwgts,
             me_prod_on = getattr(production, 'me_wgt', None)
             if not me_prod_on:
                 # The denominator has to be the same quantity as the numerator,
-                # in the same frame: Tr(rho_off) is built in the me_frame while
-                # calculate_matrix_element hands the matrix element the LAB
-                # momenta, and a helicity-restricted matrix element is not
-                # Lorentz invariant.  Unpolarised runs have no frame boost and
-                # keep the matrix-element call, bit for bit.
+                # in the same frame: Tr(rho_off) is built in the me_frame, and a
+                # helicity-restricted matrix element is not Lorentz invariant,
+                # so the lab-frame calculate_matrix_element is a different
+                # projection (on a boosted polarised event, by orders of
+                # magnitude -- which then sets the mass-stage bound).
+                # Unpolarised runs have no frame boost and keep the
+                # matrix-element call, bit for bit.
                 me_prod_on = self._onshell_production_norm(production,
                                                            prod_static)
                 production.me_wgt = me_prod_on
@@ -11614,15 +11628,21 @@ def calculate_matrix_element_from_density(self, production, decays, decay_dict,
             #VALENTIN: except for the mode "full", we should not compute the matrix element here
             MEdenom_prod, MEdenom_decay = None, None
             if not density_pole_approximation:
-                # compute the denominator and then reshuffle the event before 
-                # computing the numerator 
-                # same frame-consistency fix as on the sequential mass
-                # stage: the numerator is the contraction of the (possibly
-                # restricted) production density built in the me_frame, so the
-                # denominator cannot be the lab-frame matrix element.
+                # compute the denominator and then reshuffle the event before
+                # computing the numerator
+                #
+                # The denominator has to be taken in the same frame as the
+                # numerator (see _onshell_production_norm), as on the
+                # sequential mass stage.
+                #
+                # Not polarised-only: keep_weight_for_polarization_* or
+                # 'unweighting = auto' can bring an unbraced production here with
+                # the frame axis on. That is safe: the denominator is a constant
+                # per production event, identical across its joint trials, so it
+                # cancels out of the accept/reject.
                 MEdenom_prod = self._onshell_production_norm(production,
                                                              prod_static)
-                MEdenom_decay = 1.0              
+                MEdenom_decay = 1.0
                 for key in decays:
                     for dec in decays[key]:
                         MEdenom_decay *= self.calculate_matrix_element(dec)
@@ -11945,7 +11965,14 @@ def _frame_boost(self, event):
         if frame_id <= 0:
             return None
         _, orig_order, _, _, _ = self.get_pdir(event)
-        momenta = event.get_momenta(orig_order)
+        # merged_map: orig_order comes out of all_me, which is keyed by the
+        # MERGED-pdg tag when apply_flavor_grouping is on (81/-81 in place of
+        # u/d/...).  Without the map the raw event pdgs are looked up in a
+        # merged block and get_mapping dies with
+        # "ValueError: list.index(x): x not in list" -- exactly as the matrix
+        # element call twenty lines below already guards against.
+        momenta = event.get_momenta(orig_order,
+                                    merged_map=self._revert_merged or None)
         selected = [n for n in range(1, len(momenta) + 1) if frame_id >> n & 1]
         if not selected:
             return None
@@ -12020,8 +12047,18 @@ def get_density(self, event, position, allow_hel, ncomb, dimension,
         if orig_order is None:
             _, orig_order, _, _, tag = self.get_pdir(event)
             event._ms_orig_order_for_density = orig_order
+            # cache the tag get_pdir resolved rather than recomputing it: it is
+            # the MERGED-pdg tag (all_me is keyed by 81/-81, not by the raw
+            # event pdgs, so a bare get_tag_and_order() raises KeyError on the
+            # lookup below -- e.g. ((-1, 1), (23, 23)) against ((-81, 81),
+            # (23, 23))), and get_pdir also owns the 1 -> N antiparticle
+            # fallback, which no recomputation here would reproduce.
+            event._ms_tag_for_density = tag
         else: #in any case, we need tag to differentiate between production and decay
-            tag, _ = event.get_tag_and_order()
+            tag = getattr(event, '_ms_tag_for_density', None)
+            if tag is None:
+                tag, _ = event.get_tag_and_order(
+                    merged_particle=self._revert_merged or None)
 
 
         # Fast path: single-point momentum extraction without permutation
@@ -12036,6 +12073,11 @@ def get_density(self, event, position, allow_hel, ncomb, dimension,
             all_p = event.get_all_momenta(orig_order, merged_map=self._revert_merged or None)
             assert len(all_p) == 1, "Error: get_density can only be called for a single phase-space point"
             p = all_p[0]
+        # get_pdg below identifies the particles by EXACT momentum equality
+        # against the event record, so it has to see the lab-frame momenta.
+        # _boost_momenta builds new tuples and never touches its input, so
+        # keeping a reference is enough -- no copy needed.
+        p_lab = p
         if frame_boost is not None:
             p = self._boost_momenta(p, frame_boost, rest_leg=frame_rest_leg)
         P = rwgt_interface.ReweightInterface.invert_momenta(p)
@@ -12049,7 +12091,7 @@ def get_density(self, event, position, allow_hel, ncomb, dimension,
         need_raw_pdg = (self._revert_merged and
                         any(abs(pid) in merged_particles for pid in pdg_template))
         if need_raw_pdg:
-            pdgs = event.get_pdg(p)
+            pdgs = event.get_pdg(p_lab)
         else:
             pdgs = pdg_template
         n_changing = len(position)

tests/unit_tests/madspin/test_madspin.py

@@ -444,6 +444,10 @@ def __init__(self, frame_id, beampol, prodpol=None, vector=(), fermion=(),
         self.options['keep_weight_for_polarization_fermion'] = list(fermion)
         self.options['pure_interference'] = pure_interference
         self.model = _PIModelStub()
+        # MadSpinInterface carries this as a class attribute (it is {} unless
+        # flavour grouping is on); _frame_boost passes it to get_momenta as
+        # merged_map, so the stub has to have it too.
+        self._revert_merged = {}
         # what _production_polarization would have parsed out of the banner's
         # proc_card: {} for a brace-free production process
         self._production_polarization_cache = prodpol if prodpol else {}
@@ -473,7 +477,9 @@ class _MomentaEvent(object):
     def __init__(self, momenta):
         self.momenta = momenta
 
-    def get_momenta(self, orig_order):
+    def get_momenta(self, orig_order, merged_map=None):
+        # merged_map mirrors lhe_parser.Event.get_momenta: _frame_boost has
+        # to pass it so the ME ordering resolves under flavour grouping.
         return self.momenta
 
 
@@ -5229,6 +5235,11 @@ def __init__(self, lines):
             self.mg5cmd = self
 
         def extract_process(self, line):
+            if '[' in line:
+                # like MadSpin's tree-level mg5cmd on an NLO process line
+                raise self.InvalidCmd('Perturbation order QCD is not among the '
+                                      'perturbation orders allowed for by the '
+                                      'loop model')
             legs = []
             initial, final = line.split('>')
             for state, part in ([(False, p) for p in initial.split()] +
@@ -5266,6 +5277,20 @@ def test_partially_polarised_same_pdg(self):
         self.assertEqual(self.polarization('generate p p > z{0} z'),
                          {23: ((0,), None)})
 
+    def test_nlo_perturbation_bracket_is_ignored(self):
+        """'p p > z{0} z{0} [QCD]': the tree-level parser refuses the bracket,
+        which used to leave every NLO polarised production unrestricted (with
+        only a warning). The bracket says nothing about the braces on the legs,
+        so the restriction must be the bracket-free line's."""
+        for bracket in ('[QCD]', '[virt=QCD]', '[noborn=QCD]'):
+            self.assertEqual(
+                self.polarization('generate p p > z{0} z{0} %s' % bracket),
+                {23: ((0,),)})
+            self.assertEqual(
+                self.polarization('generate p p > w+{0} w+{T} %s' % bracket,
+                                  'add process p p > w+{0} w+{T} j %s' % bracket),
+                {24: ((0,), (-1, 1))})
+
     def test_broadcast_survives_extra_subprocesses(self):
         """The multiplicity of a broadcast pdg does not have to match between
         subprocesses -- this is the common 'generate X; add process X j' case."""
@@ -7070,6 +7095,12 @@ def _offshell_fixture(self):
         stub._slot_of = {index: slot for slot, index in enumerate(slots)}
         # |M_prod|^2 on shell, the denominator of the offshell mass-set weight
         stub.calculate_matrix_element = lambda *args, **opts: 1.0
+        # The mass stage takes that denominator through
+        # _onshell_production_norm, which returns calculate_matrix_element
+        # unchanged when there is no frame boost -- which is this stub's
+        # case, and the only one it can represent.
+        stub._onshell_production_norm = \
+            lambda production, prod_static: stub.calculate_matrix_element(production)
 
         def _no_pool(*args, **opts):
             raise TestPAUpFrontMass._NoDecayPool()