madmatrix: squared split orders, so the interference syntax stops returning the total

madgraph/iolibs/export_v4.py

@@ -221,6 +221,105 @@ def export_helas(self, HELAS_PATH):
         raise Exception("V4 model not supported by this type of exporter. Please use UFO model")
         return
 
+#===============================================================================
+# Squared split orders: which components the user's constraint keeps
+#===============================================================================
+def chosen_squared_orders(process, squared_orders):
+    """Which entries of squared_orders the user's '^2' constraint keeps.
+
+    squared_orders is the list of squared split-order tuples actually present
+    in the matrix element, as get_split_orders_mapping returns it; the process
+    carries the constraint. A False entry is a component that contributes to
+    the amplitude but that the user asked *not* to have summed into the total,
+    so a backend which cannot mask it does not compute what was asked for.
+
+    This is the boolean form of what set_chosen_SO_index writes as the Fortran
+    CHOSEN_SO_CONFIGS DATA statement, kept here so that a backend can ask the
+    question without parsing that string.
+    """
+
+    user_squared_orders = process.get('squared_orders')
+    split_orders = process.get('split_orders')
+
+    if len(user_squared_orders)==0:
+        return [True]*len(squared_orders)
+
+    res = []
+    for sqsos in squared_orders:
+        is_a_match = True
+        for user_sqso, value in user_squared_orders.items():
+            if user_sqso == 'WEIGHTED' :
+                logger.debug('WEIGHTED^2%s%s encoutered. Please check behavior for' + \
+                        'https://bazaar.launchpad.net/~maddevelopers/mg5amcnlo/3.0.1/revision/613', \
+                        (process.get_squared_order_type(user_sqso), sqsos[split_orders.index(user_sqso)]))
+            if user_sqso not in split_orders:
+                is_a_match = False
+            elif (process.get_squared_order_type(user_sqso) =='==' and \
+                    value!=sqsos[split_orders.index(user_sqso)]) or \
+               (process.get_squared_order_type(user_sqso) in ['<=','='] and \
+                            value<sqsos[split_orders.index(user_sqso)]) or \
+               (process.get_squared_order_type(user_sqso) == '>' and \
+                            value>=sqsos[split_orders.index(user_sqso)]):
+                is_a_match = False
+                break
+        res.append(is_a_match)
+
+    return res
+
+def split_order_tables(matrix_element):
+    """The squared split-order bookkeeping a backend needs, or None.
+
+    Returns a dict with
+
+      nampso      how many amplitude split orders the amplitudes fall into
+      nsqampso    how many squared orders their pairs produce
+      amp_so      {amplitude number -> amplitude-order index}, 0-based
+      sqsoindex   sqsoindex[m][n] -> squared-order index, 0-based. SYMMETRIC,
+                  because a squared order is the SUM of the two amplitude
+                  orders (Fortran SQSOINDEX), which is what lets a masked sum
+                  stay real: (m,n) and (n,m) are kept or dropped together.
+      chosen      [bool] per squared order, the user's constraint
+      names       ['QED=0', ...] per squared order, for comments
+
+    None when the process has no split orders, i.e. when there is one implicit
+    component and every backend already computes it.
+    """
+
+    process = matrix_element.get('processes')[0]
+    split_orders = process.get('split_orders')
+    if not split_orders:
+        return None
+    squared_orders, amp_orders = matrix_element.get_split_orders_mapping()
+    if not squared_orders:
+        return None
+
+    amp_so = {}
+    for iampso, (_orders, amp_numbers) in enumerate(amp_orders):
+        for namp in amp_numbers:
+            amp_so[namp] = iampso
+
+    # The squared order a pair of amplitude orders lands in: add the two
+    # amplitude orders and look the sum up. A pair whose sum is not in the
+    # list cannot happen (the list is built from exactly these sums), but
+    # guard anyway rather than write a negative index into the generated code.
+    index_of = {tuple(sqso): i for i, sqso in enumerate(squared_orders)}
+    sqsoindex = []
+    for m, (orders_m, _a) in enumerate(amp_orders):
+        row = []
+        for n, (orders_n, _b) in enumerate(amp_orders):
+            key = tuple(om + on for om, on in zip(orders_m, orders_n))
+            row.append(index_of.get(key, -1))
+        sqsoindex.append(row)
+
+    return {'nampso': len(amp_orders),
+            'nsqampso': len(squared_orders),
+            'amp_so': amp_so,
+            'sqsoindex': sqsoindex,
+            'chosen': chosen_squared_orders(process, squared_orders),
+            'names': [' '.join('%s=%d' % (o, v)
+                               for o, v in zip(split_orders, sqso))
+                      for sqso in squared_orders]}
+
 #===============================================================================
 # ProcessExporterFortran
 #===============================================================================
@@ -2017,34 +2116,9 @@ def set_chosen_SO_index(self, process, squared_orders):
         finds what indices of the squared_orders list the user intends to pick.
         It returns this as a string of comma-separated successive '.true.' or 
         '.false.' for each index."""
-
-        user_squared_orders = process.get('squared_orders')
-        split_orders = process.get('split_orders')
-
-        if len(user_squared_orders)==0:
-            return ','.join(['.true.']*len(squared_orders))
-
-        res = []
-        for sqsos in squared_orders:
-            is_a_match = True
-            for user_sqso, value in user_squared_orders.items():
-                if user_sqso == 'WEIGHTED' :
-                    logger.debug('WEIGHTED^2%s%s encoutered. Please check behavior for' + \
-                            'https://bazaar.launchpad.net/~maddevelopers/mg5amcnlo/3.0.1/revision/613', \
-                            (process.get_squared_order_type(user_sqso), sqsos[split_orders.index(user_sqso)]))
-                if user_sqso not in split_orders:
-                    is_a_match = False
-                elif (process.get_squared_order_type(user_sqso) =='==' and \
-                        value!=sqsos[split_orders.index(user_sqso)]) or \
-                   (process.get_squared_order_type(user_sqso) in ['<=','='] and \
-                                value<sqsos[split_orders.index(user_sqso)]) or \
-                   (process.get_squared_order_type(user_sqso) == '>' and \
-                                value>=sqsos[split_orders.index(user_sqso)]):
-                    is_a_match = False
-                    break
-            res.append('.true.' if is_a_match else '.false.')
-
-        return ','.join(res)
+
+        return ','.join('.true.' if keep else '.false.' for keep in
+                        chosen_squared_orders(process, squared_orders))
 
     def get_split_orders_lines(self, orders, array_name, n=5):
         """ Return the split orders definition as defined in the list orders and

madgraph/iolibs/template_files/madmatrix/color_sum_splitorders.cc

@@ -0,0 +1,153 @@
+// Copyright (C) 2020-2026 CERN and UCLouvain.
+// Licensed under the GNU Lesser General Public License (version 3 or later).
+// Created originally by: A. Valassi (Sep 2025) for the MG5aMC CUDACPP plugin.
+// Further modified by: A. Valassi (2025).
+// Integrated with the MadGraph7 project in Feb 2026.
+//
+// SQUARED SPLIT ORDERS. This is color_sum.cc for a process whose '^2'
+// constraint leaves more than one amplitude split order, and it is a separate
+// file because the sum itself changes shape rather than gaining a hole: the
+// jamps arrive as nampso vectors end to end (njampso = nampso*ncolor) and the
+// color sum pairs them, which is the Fortran GET_MATRIX contract
+// JAMP(NCOLOR,NAMPSO) -> RES(NSQAMPSO).
+//
+// THE PAIR LOOP RUNS OVER ALL nampso*nampso ORDERED PAIRS, never a triangle.
+// The color contraction below keeps the triangular matrix whose off-diagonal
+// is doubled, which for a single jamp vector is just the statement that the
+// (i,j) and (j,i) terms are conjugates. For two DIFFERENT jamp vectors they
+// are not, so the value computed for one ordered pair (m,n) is NOT that pair's
+// contribution -- only (m,n) and (n,m) together are, and they are equal to the
+// true sum of the two. That is why both orderings must be summed, and why it
+// is safe to mask on the squared order: sqSoIndex is symmetric (a squared
+// order is the SUM of the two amplitude orders), so a pair and its transpose
+// are always kept or dropped together and the masked total stays real.
+
+#include "color_sum.h"
+
+#include "mgOnGpuConfig.h"
+
+#include "MemoryAccessMatrixElements.h"
+
+#ifdef MGONGPUCPP_GPUIMPL
+#error The squared split-order color sum is implemented for the CPU (SIMD) backend only. \
+The GPU jamp buffers are sized for a single jamp vector per helicity (ncolor, not njampso), \
+so a GPU build of this process would silently sum the wrong thing. Use a CPU backend, or \
+generate the process with a constraint that leaves a single amplitude split order.
+#endif
+
+namespace mg5amcCpu
+{
+  constexpr int ncolor = CPPProcess::ncolor;     // the number of leading colors
+  constexpr int nampso = CPPProcess::nampso;     // the amplitude split orders
+  constexpr int njampso = CPPProcess::njampso;   // ncolor * nampso: the jamps of every order, end to end
+  constexpr int nsqampso = CPPProcess::nsqampso; // the squared orders their pairs produce
+
+  //--------------------------------------------------------------------------
+
+  // *** COLOR MATRIX BELOW ***
+%(color_matrix_lines)s
+
+  //--------------------------------------------------------------------------
+
+  // *** SQUARED SPLIT ORDERS BELOW ***
+%(sqso_tables)s
+
+  //--------------------------------------------------------------------------
+
+  void
+  color_sum_cpu( fptype* allMEs,              // output: allMEs[nevt], add |M|^2 for one specific helicity
+                 const cxtype_sv* allJamp_sv, // input: jamp_sv[njampso] (float/double) or jamp_sv[2*njampso] (mixed) for one specific helicity
+                 const int ievt0 )            // input: first event number in current C++ event page (for CUDA, ievt depends on threadid)
+  {
+    // Pre-compute a constexpr triangular color matrix properly normalized #475
+    struct TriangularNormalizedColorMatrix
+    {
+      // See https://stackoverflow.com/a/34465458
+      __host__ __device__ constexpr TriangularNormalizedColorMatrix()
+        : value()
+      {
+        for( int icol = 0; icol < ncolor; icol++ )
+        {
+          // Diagonal terms
+          value[icol][icol] = colorMatrix[icol][icol] / colorDenom[icol];
+          // Off-diagonal terms
+          for( int jcol = icol + 1; jcol < ncolor; jcol++ )
+            value[icol][jcol] = 2 * colorMatrix[icol][jcol] / colorDenom[icol];
+        }
+      }
+      fptype2 value[ncolor][ncolor];
+    };
+    static constexpr auto cf2 = TriangularNormalizedColorMatrix();
+    fptype_sv deltaMEs = { 0 };
+#if defined MGONGPU_CPPSIMD and defined MGONGPU_FPTYPE_DOUBLE and defined MGONGPU_FPTYPE2_FLOAT
+    fptype_sv deltaMEs_next = { 0 };
+#endif
+    // Gather the color flows the sum runs over, for each amplitude split order.
+    // The order index strides by ncolor, exactly as calculate_jamps wrote them.
+    // NB in mixed mode the two neppV vectors of allJamp_sv, at ijamp and at
+    // njampso+ijamp, are two halves of the event page and not two colors.
+    fptype2_sv jampR_sv[nampso][ncolor];
+    fptype2_sv jampI_sv[nampso][ncolor];
+    for( int iao = 0; iao < nampso; iao++ )
+    {
+      for( int icol = 0; icol < ncolor; icol++ )
+      {
+        const int ijamp = iao * ncolor + icol;
+#if defined MGONGPU_CPPSIMD and defined MGONGPU_FPTYPE_DOUBLE and defined MGONGPU_FPTYPE2_FLOAT
+        jampR_sv[iao][icol] = fpvmerge( cxreal( allJamp_sv[ijamp] ), cxreal( allJamp_sv[njampso + ijamp] ) );
+        jampI_sv[iao][icol] = fpvmerge( cximag( allJamp_sv[ijamp] ), cximag( allJamp_sv[njampso + ijamp] ) );
+#else
+        jampR_sv[iao][icol] = (fptype2_sv)( cxreal( allJamp_sv[ijamp] ) );
+        jampI_sv[iao][icol] = (fptype2_sv)( cximag( allJamp_sv[ijamp] ) );
+#endif
+      }
+    }
+    // Loop over ORDERED pairs of amplitude split orders (see the note at the top:
+    // this must not be turned into a triangle over iao/jao)
+    for( int iao = 0; iao < nampso; iao++ )
+    {
+      for( int jao = 0; jao < nampso; jao++ )
+      {
+        // The squared order this pair contributes to, dropped if the process
+        // asked for a contribution that does not include it
+        if( !chosenSqso[sqSoIndex[iao][jao]] ) continue;
+        // Loop over icol
+        for( int icol = 0; icol < ncolor; icol++ )
+        {
+          // Diagonal terms (the color contraction is still triangular)
+          fptype2_sv ztempR_sv = cf2.value[icol][icol] * jampR_sv[iao][icol];
+          fptype2_sv ztempI_sv = cf2.value[icol][icol] * jampI_sv[iao][icol];
+          // Loop over jcol
+          for( int jcol = icol + 1; jcol < ncolor; jcol++ )
+          {
+            // Off-diagonal terms
+            ztempR_sv += cf2.value[icol][jcol] * jampR_sv[iao][jcol];
+            ztempI_sv += cf2.value[icol][jcol] * jampI_sv[iao][jcol];
+          }
+          // ztemp comes from the jamps of order iao, and is contracted with those
+          // of order jao (Fortran: ZTEMP from JAMP(:,M), times DCONJG(JAMP(I,N)))
+          fptype2_sv deltaMEs2 = ( jampR_sv[jao][icol] * ztempR_sv + jampI_sv[jao][icol] * ztempI_sv ); // may underflow #831
+#if defined MGONGPU_CPPSIMD and defined MGONGPU_FPTYPE_DOUBLE and defined MGONGPU_FPTYPE2_FLOAT
+          deltaMEs += fpvsplit0( deltaMEs2 );
+          deltaMEs_next += fpvsplit1( deltaMEs2 );
+#else
+          deltaMEs += deltaMEs2;
+#endif
+        }
+      }
+    }
+    // *** STORE THE RESULTS ***
+    using E_ACCESS = HostAccessMatrixElements; // non-trivial access: buffer includes all events
+    fptype* MEs = E_ACCESS::ieventAccessRecord( allMEs, ievt0 );
+    // NB: color_sum ADDS |M|^2 for one helicity to the running sum of |M|^2 over helicities for the given event(s)
+    fptype_sv& MEs_sv = E_ACCESS::kernelAccess( MEs );
+    MEs_sv += deltaMEs; // fix #435
+#if defined MGONGPU_CPPSIMD and defined MGONGPU_FPTYPE_DOUBLE and defined MGONGPU_FPTYPE2_FLOAT
+    fptype* MEs_next = E_ACCESS::ieventAccessRecord( allMEs, ievt0 + neppV );
+    fptype_sv& MEs_sv_next = E_ACCESS::kernelAccess( MEs_next );
+    MEs_sv_next += deltaMEs_next;
+#endif
+  }
+
+  //--------------------------------------------------------------------------
+}

madgraph/iolibs/template_files/madmatrix/process_class.inc

@@ -58,7 +58,7 @@
     static constexpr int ncolor = %(ncolor)s; // the number of leading colors: e.g. 1 for e+ e- -> mu+ mu-
     // The color structures the color flow is picked among. Same as ncolor,
     // unless the color sum runs on a smaller basis (see set_color_flow_lines_cpp).
-    static constexpr int ncolor_flow = %(ncolor_flow)s;
+    static constexpr int ncolor_flow = %(ncolor_flow)s;%(split_order_constants)s
     static constexpr int nmaxflavor = %(nmaxflavor)d; // the maximum number of flavor combinations
 
     // Hardcoded parameters for this process (constant class variables)

madgraph/iolibs/template_files/madmatrix/process_function_definitions.inc

@@ -84,7 +84,7 @@ namespace mg5amcCpu
   constexpr int ncomb = CPPProcess::ncomb;   // #helicity combinations: e.g. 16 for e+ e- -> mu+ mu- (2**4 = fermion spin up/down ** npar)
   constexpr int ncolor = CPPProcess::ncolor; // the number of leading colors
   constexpr int ncolor_flow = CPPProcess::ncolor_flow; // the color structures the color flow is picked among
-  constexpr int nmaxflavor = CPPProcess::nmaxflavor; // the maximum number of flavor combinations
+  constexpr int nmaxflavor = CPPProcess::nmaxflavor; // the maximum number of flavor combinations%(jampso_aliases)s
 
   // [NB: I am currently unable to get the right value of nwf in CPPProcess.h - will hardcode it in CPPProcess.cc instead (#644)]
   //using CPPProcess::nwf; // #wavefunctions = #external (npar) + #internal: e.g. 5 for e+ e- -> mu+ mu- (1 internal is gamma or Z)
@@ -576,9 +576,9 @@ namespace mg5amcCpu
         constexpr fptype_sv* jamp2_sv = nullptr; // no need for color selection during helicity filtering
         //std::cout << "sigmaKin_getGoodHel ihel=" << ihel << ( isGoodHel[ihel] ? " true" : " false" ) << std::endl;
 #if defined MGONGPU_CPPSIMD and defined MGONGPU_FPTYPE_DOUBLE and defined MGONGPU_FPTYPE2_FLOAT
-        cxtype_sv jamp_sv[2 * ncolor] = {}; // all zeros
+        cxtype_sv jamp_sv[2 * %(jamp_ncolor)s] = {}; // all zeros
 #else
-        cxtype_sv jamp_sv[ncolor] = {};  // all zeros
+        cxtype_sv jamp_sv[%(jamp_ncolor)s] = {};  // all zeros
 #endif
         calculate_jamps( ihel, allmomenta, allcouplings, hgFlavorVec, jamp_sv, false, allNumerators, allDenominators, jamp2_sv, ievt00 ); //maxtry?
         color_sum_cpu( allMEs, jamp_sv, ievt00 );

madgraph/iolibs/template_files/madmatrix/process_matrix.inc

@@ -16,7 +16,7 @@
       if( jamp2_sv ) // disable color choice if nullptr
       {
         for( int icol = 0; icol < ncolor_flow; icol++ )
-          jamp2_sv[ncolor_flow * iParity + icol] += cxabs2( %(jamp_flow)s[icol] ); // may underflow #831
+          jamp2_sv[ncolor_flow * iParity + icol] += cxabs2( %(jamp_flow_col)s ); // may underflow #831
       }
 #else /* clang-format off */
       assert( iParity == 0 ); // sanity check for J2_ACCESS
@@ -25,7 +25,7 @@
       {
         for( int icol = 0; icol < ncolor_flow; icol++ )
           // NB: atomicAdd is needed after moving to cuda streams with one helicity per stream!
-          atomicAdd( &J2_ACCESS::kernelAccessIcol( colAllJamp2s, icol ), cxabs2( %(jamp_flow)s[icol] ) );
+          atomicAdd( &J2_ACCESS::kernelAccessIcol( colAllJamp2s, icol ), cxabs2( %(jamp_flow_col)s ) );
       }
 #endif /* clang-format on */
 
@@ -35,12 +35,12 @@
       // In CUDA, copy the local jamp to the output global-memory jamp
       constexpr int ihel0 = 0; // the allJamps buffer already points to a specific helicity _within a super-buffer for dcNGoodHel helicities_
       using J_ACCESS = DeviceAccessJamp;
-      for( int icol = 0; icol < ncolor; icol++ )
+      for( int icol = 0; icol < %(jamp_ncolor)s; icol++ )
         J_ACCESS::kernelAccessIcolIhelNhel( allJamps, icol, ihel0, dcNGoodHel ) = jamp_sv[icol];
 #else
       // In C++, copy the local jamp to the output array passed as function argument
-      for( int icol = 0; icol < ncolor; icol++ )
-        allJamp_sv[iParity * ncolor + icol] = jamp_sv[icol];
+      for( int icol = 0; icol < %(jamp_ncolor)s; icol++ )
+        allJamp_sv[iParity * %(jamp_ncolor)s + icol] = jamp_sv[icol];
 #endif
     }
     // END LOOP ON IPARITY

madgraph/iolibs/template_files/madmatrix/process_sigmaKin_function.inc

@@ -122,7 +122,7 @@
 %(cpp_blas_helicity_loop)s      for( int ighel = 0; ighel < cNGoodHel; ighel++ )
       {
         const int ihel = cGoodHel[ighel];
-        cxtype_sv jamp_sv[nParity * ncolor] = {}; // fixed nasty bug (omitting 'nParity' caused memory corruptions after calling calculate_jamps)
+        cxtype_sv jamp_sv[nParity * %(jamp_ncolor)s] = {}; // fixed nasty bug (omitting 'nParity' caused memory corruptions after calling calculate_jamps)
         // **NB! in "mixed" precision, using SIMD, calculate_jamps computes MEs for TWO neppV pages with a single channelId! #924
         bool storeChannelWeights = allChannelIds != nullptr || allrnddiagram != nullptr;
         calculate_jamps( ihel, allmomenta, allcouplings, iflavorVec, jamp_sv, storeChannelWeights, allNumerators, allDenominators, jamp2_sv, ievt00 );