diff --git a/madgraph/iolibs/export_v4.py b/madgraph/iolibs/export_v4.py index 49a59dff11..75f94213dd 100755 --- a/madgraph/iolibs/export_v4.py +++ b/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' 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' 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 diff --git a/madgraph/iolibs/template_files/madmatrix/color_sum_splitorders.cc b/madgraph/iolibs/template_files/madmatrix/color_sum_splitorders.cc new file mode 100644 index 0000000000..f9f2fcb7fc --- /dev/null +++ b/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 + } + + //-------------------------------------------------------------------------- +} diff --git a/madgraph/iolibs/template_files/madmatrix/process_class.inc b/madgraph/iolibs/template_files/madmatrix/process_class.inc index 87c6242721..501a2c3793 100644 --- a/madgraph/iolibs/template_files/madmatrix/process_class.inc +++ b/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) diff --git a/madgraph/iolibs/template_files/madmatrix/process_function_definitions.inc b/madgraph/iolibs/template_files/madmatrix/process_function_definitions.inc index c115970337..c6db8e9ef1 100644 --- a/madgraph/iolibs/template_files/madmatrix/process_function_definitions.inc +++ b/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 ); diff --git a/madgraph/iolibs/template_files/madmatrix/process_matrix.inc b/madgraph/iolibs/template_files/madmatrix/process_matrix.inc index f266f9b198..03aaa94d36 100644 --- a/madgraph/iolibs/template_files/madmatrix/process_matrix.inc +++ b/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 diff --git a/madgraph/iolibs/template_files/madmatrix/process_sigmaKin_function.inc b/madgraph/iolibs/template_files/madmatrix/process_sigmaKin_function.inc index ee74555d10..de51e2b556 100644 --- a/madgraph/iolibs/template_files/madmatrix/process_sigmaKin_function.inc +++ b/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 ); diff --git a/madmatrix/model_handling.py b/madmatrix/model_handling.py index 5b5581fcb0..e41814347d 100644 --- a/madmatrix/model_handling.py +++ b/madmatrix/model_handling.py @@ -1677,6 +1677,17 @@ def get_process_class_definitions(self, write=True): replace_dict['nbhel'] = self.matrix_elements[0].get_helicity_combinations() # number of helicity combinations replace_dict['ndiagrams'] = len(self.matrix_elements[0].get('diagrams')) # AV FIXME #910: elsewhere matrix_element.get('diagrams') and max(config[0]... replace_dict['nmaxflavor'] = len(self.matrix_elements[0].get_external_flavors_with_iden()) # number of flavor combinations + # Only written when the jamps are actually split, so that a process + # without squared split orders keeps the header it always had + so = self.split_orders_info() + replace_dict['split_order_constants'] = '' if not self.split_orders_active() else ( + '\n // Squared split orders: the amplitudes fall into nampso amplitude' + '\n // orders, the jamps carry one vector per order (njampso long in total)' + '\n // and the color sum pairs them into nsqampso squared orders' + '\n // (see color_sum.cc, written from color_sum_splitorders.cc).' + '\n static constexpr int nampso = %d;' + '\n static constexpr int njampso = ncolor * nampso; // the jamps of every amplitude order, end to end' + '\n static constexpr int nsqampso = %d;' % (so['nampso'], so['nsqampso'])) replace_dict['nwave'] = 4 if (fd_gauge): replace_dict['nwave'] += 1 @@ -1699,6 +1710,12 @@ def get_process_function_definitions(self, write=True): """The complete class definition for the process""" replace_dict = super().get_process_function_definitions(write=False) # defines replace_dict['initProc_lines'] replace_dict['hardcoded_initProc_lines'] = replace_dict['initProc_lines'].replace( 'm_pars->', 'Parameters::') + replace_dict['jamp_ncolor'] = self.jamp_ncolor() + # Only pulled into scope when the jamps are split, so that a process + # without split orders keeps exactly the constants it always had + replace_dict['jampso_aliases'] = '' if not self.split_orders_active() else ( + '\n constexpr int nampso = CPPProcess::nampso; // the amplitude split orders' + '\n constexpr int njampso = CPPProcess::njampso; // ncolor * nampso: the jamps of every order, end to end') couplings2order_indep = [] ###replace_dict['ncouplings'] = len(self.couplings2order) ###replace_dict['ncouplingstimes2'] = 2 * replace_dict['ncouplings'] @@ -1924,6 +1941,7 @@ def get_sigmaKin_lines(self, color_amplitudes, write=True): replace_dict = super().get_sigmaKin_lines(color_amplitudes, write=False) replace_dict['proc_id'] = self.proc_id if self.proc_id>0 else 1 replace_dict['proc_id_source'] = 'MadMatrix exporter' + replace_dict['jamp_ncolor'] = self.jamp_ncolor() # Extract denominator (avoid to extend size for mirroring) den_factors = [str(me.get_denominator_factor()) for me in \ @@ -1961,6 +1979,9 @@ def get_sigmaKin_lines(self, color_amplitudes, write=True): def get_all_sigmaKin_lines(self, color_amplitudes, class_name): """Get sigmaKin_process for all subprocesses for CPPProcess.cc""" ret_lines = [] + # The jamps are one vector per amplitude split order, njampso long in + # total; 'ncolor' without them, so the default output is unchanged. + jamp_dim = 'njampso' if self.split_orders_active() else 'ncolor' if self.single_helicities: ###misc.sprint(type(self.helas_call_writer)) ###misc.sprint( 'before get_matrix_element_calls', self.matrix_elements[0].get_number_of_wavefunctions() ) # WRONG value of nwf, eg 7 for gg_tt @@ -2093,7 +2114,7 @@ def get_all_sigmaKin_lines(self, color_amplitudes, class_name): ret_lines.append(""" // Local variables for the given CUDA event (ievt) or C++ event page (ipagV) // [jamp: sum (for one event or event page) of the invariant amplitudes for all Feynman diagrams in a given color combination] - cxtype_sv jamp_sv[ncolor] = {}; // all zeros (NB: vector cxtype_v IS initialized to 0, but scalar cxtype is NOT, if "= {}" is missing!)""") + cxtype_sv jamp_sv[%s] = {}; // all zeros (NB: vector cxtype_v IS initialized to 0, but scalar cxtype is NOT, if "= {}" is missing!)""" % jamp_dim) # Shared sub-expressions of the color flows, filled in while the # amplitudes go by (see MadMatrixUFOHelasCallWriter.build_jamp_plan). # No "= {}": each one is assigned before it is ever read. @@ -2232,6 +2253,51 @@ def blas_available_flags(cls): return '' return cls._blas_flags + # AV - new method (add the split-order holes to process_matrix.inc) + def get_matrix_single_process(self, i, matrix_element, color_amplitudes, + class_name, write=True): + replace_dict = super().get_matrix_single_process( + i, matrix_element, color_amplitudes, class_name, write=False) + replace_dict['jamp_ncolor'] = self.jamp_ncolor() + # set_color_flow_lines_cpp fills jamp_flow / jamp_flow_col; it runs from + # get_process_class_definitions, before this, but be explicit rather + # than rely on the ordering of two independent methods. + if 'jamp_flow_col' not in replace_dict: + self.set_color_flow_lines_cpp(matrix_element, replace_dict) + if write: + return self.read_template_file(self.single_process_template) % replace_dict + return replace_dict + + # AV - new method + def jamp_ncolor(self): + """The length of a jamp array: 'ncolor', or 'njampso' (= ncolor*nampso) + once the jamps carry an amplitude-order index. Templates spell the size + through this hole so that a process without split orders gets exactly + the text it got before they existed.""" + return 'njampso' if self.split_orders_active() else 'ncolor' + + # AV - new method + def split_orders_info(self): + """The squared split-order tables for this process, or None. + + None means the process has no '^2' constraint, and then every hole this + fills reproduces the code that was written before split orders existed + -- which is how the default path stays byte-for-byte what it was. + """ + if not hasattr(self, '_split_orders_info'): + from madgraph.iolibs.export_v4 import split_order_tables + self._split_orders_info = split_order_tables(self.matrix_elements[0]) + return self._split_orders_info + + def split_orders_active(self): + """Whether the jamps carry an amplitude-order index at all. + + A single amplitude order is the same code as none: the pair loop of the + color sum has one term, so it is left switched off rather than writing + a 1x1 interference matrix.""" + so = self.split_orders_info() + return bool(so) and so['nampso'] > 1 + # AV - new method @classmethod def cpp_blas_wanted_for(cls, ncolor): @@ -2245,17 +2311,53 @@ def cpp_blas_wanted_for(cls, ncolor): return ncolor >= cls.blas_min_ncolor def cpp_blas_wanted(self): + # The BLAS color sum multiplies a single jamp vector per helicity by the + # color matrix. With the jamps split by amplitude order the sum is a + # pair loop instead, so the scalar color sum is used for those. + if self.split_orders_active(): + return False return self.cpp_blas_wanted_for( max(1, len(self.matrix_elements[0].get('color_basis')))) + # AV - new method + def get_sqso_table_lines(self): + """The interference matrix and the squared-order mask, as C++ tables. + + sqSoIndex[m][n] is the Fortran SQSOINDEX: the squared order a pair of + amplitude orders lands in. It is symmetric because a squared order is + the SUM of the two amplitude orders, which is what lets color_sum_cpu + mask on it and still get a real answer.""" + so = self.split_orders_info() + lines = [] + lines.append(' // The squared order each pair of amplitude orders' + ' contributes to (symmetric)') + lines.append(' static constexpr int sqSoIndex[nampso][nampso] = {') + lines.append(',\n'.join(' { %s }' % ', '.join(str(i) for i in row) + for row in so['sqsoindex'])) + lines.append(' };') + lines.append(' // The squared orders the process asked for:') + for k, (name, keep) in enumerate(zip(so['names'], so['chosen'])): + lines.append(' // %d) %s%s' % (k, name, + '' if keep else ' [dropped]')) + lines.append(' static constexpr bool chosenSqso[nsqampso] = { %s };' + % ', '.join('true' if k else 'false' + for k in so['chosen'])) + return '\n'.join(lines) + # AV - new method def edit_colorsum(self): """Generate color_sum.cc""" ###misc.sprint('Entering OneProcessExporterMadMatrix.edit_colorsum') - template = open(pjoin(self.template_path,'madmatrix','color_sum.cc'),'r').read() + # A process whose '^2' constraint leaves more than one amplitude split + # order gets the dedicated pair-loop color sum instead (see that file). + split = self.split_orders_active() + name = 'color_sum_splitorders.cc' if split else 'color_sum.cc' + template = open(pjoin(self.template_path,'madmatrix',name),'r').read() replace_dict = {} # Extract color matrix again (this was also in get_matrix_single_process called within get_all_sigmaKin_lines) replace_dict['color_matrix_lines'] = self.get_color_matrix_lines(self.matrix_elements[0]) + if split: + replace_dict['sqso_tables'] = self.get_sqso_table_lines() replace_dict['cpp_blas_color_sum'] = '' if self.cpp_blas_wanted(): replace_dict['cpp_blas_color_sum'] = strip_banner( @@ -2412,23 +2514,50 @@ def set_color_flow_lines_cpp(self, matrix_element, replace_dict): color_basis = matrix_element.get('color_basis') flow_basis = color_basis.get_flow_basis() if color_basis else None + # A color flow is picked from the WHOLE amplitude, so when the jamps are + # split by amplitude order the flow is built from their sum: the color + # flow amplitude of flow i is sum_M jamp_M[i], and squaring that gives + # back exactly today's jamp2 when there is a single amplitude order. + # (The squared-order mask is deliberately not applied here: this weight + # only chooses which color flow to write out, it is not an observable.) + so = self.split_orders_info() + split = bool(so) and so['nampso'] > 1 + jamp_src = 'jamp_sv' + so_lines = [] + if split: + jamp_src = 'jampso_sv' + so_lines = ['', + ' // The color flows of the whole amplitude: the' + ' jamps of every', + ' // amplitude split order summed back together' + ' (see color_sum.cc,', + ' // which instead keeps them apart and pairs' + ' them)', + ' cxtype_sv jampso_sv[ncolor] = {};', + ' for( int icol = 0; icol < ncolor; icol++ )', + ' for( int iao = 0; iao < nampso; iao++ )', + ' jampso_sv[icol] += jamp_sv[iao * ncolor +' + ' icol];'] + if flow_basis is None or flow_basis is color_basis: replace_dict['ncolor_flow'] = replace_dict['ncolor'] - replace_dict['jampflow_lines'] = '' - replace_dict['jamp_flow'] = 'jamp_sv' + replace_dict['jampflow_lines'] = '\n'.join(so_lines) + replace_dict['jamp_flow'] = jamp_src + replace_dict['jamp_flow_col'] = '%s[icol]' % jamp_src return projection = color_basis.get_flow_projection() - lines = ['', + lines = so_lines + ['', ' // The color flow jamps, rebuilt from the ones entering', ' // the color sum through the Kleiss-Kuijf relations', ' cxtype_sv jampf_sv[ncolor_flow] = {};'] for i, coeff_list in enumerate(projection): - terms = ''.join('%sjamp_sv[%d]' % (self.coeff(coefficient[0], - coefficient[1], - coefficient[2], - coefficient[3]), - number - 1) + terms = ''.join('%s%s[%d]' % (self.coeff(coefficient[0], + coefficient[1], + coefficient[2], + coefficient[3]), + jamp_src, + number - 1) for coefficient, number in coeff_list) lines.append(' jampf_sv[%d] = %s;' % (i, terms if terms else 'cxzero_sv()')) @@ -2436,6 +2565,7 @@ def set_color_flow_lines_cpp(self, matrix_element, replace_dict): replace_dict['ncolor_flow'] = max(1, len(flow_basis)) replace_dict['jampflow_lines'] = '\n'.join(lines) replace_dict['jamp_flow'] = 'jampf_sv' + replace_dict['jamp_flow_col'] = 'jampf_sv[icol]' logger.debug('Color sum on %d DDM structures, color flow on %d trace ' 'structures (%d Kleiss-Kuijf terms)', @@ -2814,6 +2944,20 @@ def jamp_statement(cls, target, terms, assign): pieces.append('%s %s%s' % ('-' if sign < 0 else '+', factor, name)) return '%s %s %s;' % (target, '=' if assign else '+=', ' '.join(pieces)) + def split_order_index(self, matrix_element): + """{amplitude number -> amplitude-order index}, or None. + + None whenever the jamps are a single vector, i.e. no '^2' constraint or + only one amplitude order, in which case every jamp index below is the + plain color index it always was.""" + key = id(matrix_element) + if getattr(self, '_so_key', None) != key: + from madgraph.iolibs.export_v4 import split_order_tables + so = split_order_tables(matrix_element) + self._so_key = key + self._so_index = so['amp_so'] if so and so['nampso'] > 1 else None + return self._so_index + def build_jamp_plan(self, matrix_element, color_amplitudes): """Work out how the color flows are built from shared sub-expressions, and return (ntmp, captures, combines, final): @@ -2828,6 +2972,13 @@ def build_jamp_plan(self, matrix_element, color_amplitudes): if not self.jamp_optim_enabled(): return None + # A shared sub-expression is a partial sum of amplitudes. With the jamps + # split by amplitude order, two amplitudes of different orders must not + # end up in the same partial sum -- it would be added to one jamp vector + # and the order information lost. Write the flows out one amplitude at a + # time instead. + if self.split_order_index(matrix_element) is not None: + return None all_element = self.jamp_matrix(color_amplitudes) if not all_element: return None @@ -2921,6 +3072,11 @@ def super_get_matrix_element_calls(self, matrix_element, color_amplitudes, multi # one (color flow, amplitude) pair at a time, as before) jamp_plan = self.build_jamp_plan(matrix_element, color_amplitudes) self.nb_tmp_jamp = jamp_plan[0] if jamp_plan else 0 + so_index = self.split_order_index(matrix_element) + ncolor_jamp = len(color_amplitudes) + # 'ncolor' unless the jamps carry an amplitude-order index, so that a + # process without split orders gets exactly the text it always got + jamp_dim = 'njampso' if so_index is not None else 'ncolor' if jamp_plan is not None: _ntmp, jamp_captures, jamp_combines, jamp_final = jamp_plan me = matrix_element.get('diagrams') @@ -2991,7 +3147,7 @@ def super_get_matrix_element_calls(self, matrix_element, color_amplitudes, multi FLV_COUPLING_ARRAY flvCOUPs_dep{ cDPF_partner1, cDPF_partner2, dpf_value }; // Reset color flows (reset jamp_sv) at the beginning of a new event or event page - for( int i = 0; i < ncolor; i++ ) { jamp_sv[i] = cxzero_sv(); } + for( int i = 0; i < """ + jamp_dim + """; i++ ) { jamp_sv[i] = cxzero_sv(); } // Numerators for the current event (CUDA) or SIMD event page (C++) // (denominators are no longer accumulated here: they are derived as the sum of numerators later) @@ -3110,6 +3266,13 @@ def _guard_open(group_mask): gmask = diag_group_mask.get(id(diagram)) before_guard = [] if jamp_plan is None: + # With split orders the jamps are nampso vectors end to end + # and this amplitude belongs to exactly one of them, so its + # color flows are offset onto that vector. so_index is None + # otherwise and the index is the plain color index. + jamp_offset = 0 + if so_index is not None: + jamp_offset = so_index.get(namp, 0) * ncolor_jamp for njamp, coeff in color[namp].items(): scoeff = OneProcessExporterMadMatrix.coeff(*coeff) # AV if scoeff[0] == '+' : scoeff = scoeff[1:] @@ -3118,6 +3281,7 @@ def _guard_open(group_mask): scoeff = scoeff.replace(',',', ') scoeff = scoeff.replace('*',' * ') scoeff = scoeff.replace('/',' / ') + njamp = njamp + jamp_offset if scoeff.startswith('-'): amp_block.append('jamp_sv[%s] -= %samp_sv[0];' % (njamp, scoeff[1:])) # AV else: amp_block.append('jamp_sv[%s] += %samp_sv[0];' % (njamp, scoeff)) # AV else: diff --git a/madmatrix/output.py b/madmatrix/output.py index c72c7f8759..d9773eef08 100644 --- a/madmatrix/output.py +++ b/madmatrix/output.py @@ -54,6 +54,10 @@ class ProcessExporterMadMatrix(export_cpp.ProcessExporterMG7): # AV - keep OM's default for this plugin (using grouped_mode=False, "can decide to merge uu~ and u~u anyway") sa_symmetry = True + # The name this exporter is reached by on the 'output' line, for the error + # messages that have to name it back to the user. + format_name = 'mg7' + # The color sum can run on the (n-2)! Del Duca-Dixon-Maltoni basis for a # multi-gluon process, but a color flow still has to be picked among the # (n-1)! trace structures, so the trace basis is built alongside and the @@ -206,8 +210,46 @@ def write_p_makefiles(self, model): rendered = self.read_template_file(pjoin(self.madmatrix_templates, name)) % replace_dict open(pjoin(self.dir_path, 'SubProcesses', name), 'w').write(rendered) + def check_split_orders(self, matrix_element): + """Report what a squared-order constraint will produce here. + + Supported: the jamps carry an amplitude-order index and the color sum + pairs them (color_sum_splitorders.cc, the Fortran GET_MATRIX contract), + so a '^2' constraint that keeps only some squared orders gets the + contribution it asked for rather than the total. That is what makes the + interference case work -- `u u~ > t t~ QED^2==2` keeps all three + diagrams and wants the QCD-EW cross term alone, which no amount of + dropping diagrams at generation can produce. + + Not supported: a GPU build of such a process. The device jamp buffers + are sized for one jamp vector per helicity (ncolor, not njampso), and + the backend is a make-time choice rather than an output-time one, so + the refusal cannot live here: color_sum_splitorders.cc #errors under + MGONGPUCPP_GPUIMPL instead. Say so now rather than let a GPU build be + the first the user hears of it. + """ + + so = export_v4.split_order_tables(matrix_element) + if not so or so['nampso'] <= 1: + return + process = matrix_element.get('processes')[0] + kept = [n for n, k in zip(so['names'], so['chosen']) if k] + dropped = [n for n, k in zip(so['names'], so['chosen']) if not k] + logger.info( + "%s: '%s' has %d squared-order components (%s); keeping %s%s. " + "The jamps are split over %d amplitude orders and the color sum " + "pairs them; CPU backends only (a GPU build of this process will " + "not compile, by design).", + self.__class__.format_name, + process.nice_string().replace('Process: ', ''), + so['nsqampso'], ', '.join(so['names']), + ', '.join(kept) if kept else 'nothing', + '' if not dropped else ', dropping %s' % ', '.join(dropped), + so['nampso']) + # AV - add debug printouts (in addition to the default one from OM's tutorial) def generate_subprocess_directory(self, matrix_element, cpp_helas_call_writer, proc_number=None): + self.check_split_orders(matrix_element) # Propagate the --mask toggle to the helas call writer that emits the # guarded wavefunction/amplitude calls, and the output command line as # a whole for the --jamp_optim toggle of the color-flow optimisation. @@ -240,6 +282,8 @@ def modify_grouping(self, matrix_element): # so that when running `make` in a P* folder, it builds check_sa.exe as well as the process library (predicatable behaviour) class ProcessExporterMadMatrixStandalone(ProcessExporterMadMatrix): + format_name = 'standalone' + # Each P* directory links madmatrix_standalone.mk (which itself includes # madmatrix.mk) as its 'makefile'; both have to be rendered in SubProcesses/ p_makefiles = ProcessExporterMadMatrix.p_makefiles + ['madmatrix_standalone.mk'] diff --git a/tests/acceptance_tests/test_cmd.py b/tests/acceptance_tests/test_cmd.py index 5100f7c763..bbe70b49d6 100755 --- a/tests/acceptance_tests/test_cmd.py +++ b/tests/acceptance_tests/test_cmd.py @@ -1334,6 +1334,68 @@ def get_values(output_format, check_exe, build_source=False): 'all matrix elements vanished for u u~ > j j') self._assert_me_lists_close(mg7, standalone, atol=1e-7) + def test_standalone_split_orders_interference(self): + """standalone (madmatrix) must return the squared-order contribution asked for. + + The madmatrix jamps carry an amplitude-order index and the color sum + pairs them, so a '^2' constraint that keeps only some of the squared + orders gets that contribution and not the total. The case that matters + is an interference term, which cannot be reached by dropping diagrams + at generation: ``u u~ > u u~ QED^2==2`` keeps every diagram and wants + the QCD-EW cross term alone, -5.5828746494657265e-02 from the Fortran + split-order driver, where a backend with no mask returns the whole + +2.7756451199752394. + + The three components are checked to sum back to the unconstrained + total *as computed by this same backend*. That comparison is the one + that pins the pair loop: it uses one set of momenta and one set of + parameters, so it is sensitive to the interference algebra alone -- + in particular to the fact that the color contraction keeps its + doubled triangle while the loop over amplitude-order pairs must run + over all ordered pairs, since for two different jamp vectors a pair + and its transpose are not each other's conjugate. + """ + devnull = open(os.devnull, 'w') + me_re = re.compile(r'Matrix element\s*=\s*([\d.eE+-]+)\s*GeV', + re.IGNORECASE) + + def value(constraint): + if os.path.isdir(self.out_dir): + shutil.rmtree(self.out_dir) + self.do('generate u u~ > u u~ %s' % constraint) + self.do('output standalone %s -f' % self.out_dir) + proc_root = pjoin(self.out_dir, 'SubProcesses') + dirs = [d for d in os.listdir(proc_root) + if d.startswith('P') and os.path.isdir(pjoin(proc_root, d))] + self.assertTrue(dirs, 'no subprocess for %s' % constraint) + proc_dir = pjoin(proc_root, dirs[0]) + # FPTYPE=d: mixed precision hides and fakes differences here + self.assertEqual(0, subprocess.call(['make', 'FPTYPE=d'], + stdout=devnull, stderr=devnull, + cwd=proc_dir), + 'standalone %s did not build' % constraint) + log = pjoin(proc_dir, 'check.log') + subprocess.call('./check_sa.exe 1000', shell=True, cwd=proc_dir, + stdout=open(log, 'w'), stderr=subprocess.STDOUT) + found = me_re.findall(open(log).read()) + self.assertTrue(found, 'no matrix element (see %s)' % log) + return float(found[0]) + + self.do('import model sm') + interference = value('QED^2==2') + # The Fortran split-order component, to the tolerance this backend is + # compared at elsewhere (the EW couplings differ in the last digits) + self.assertAlmostEqual(interference, -5.5828746494657265e-02, delta=1e-7) + # ... and emphatically not the unmasked total + self.assertLess(abs(interference), 1.0) + + components = [value('QED^2==0'), interference, value('QED^2==4')] + total = value('QED^2<=4') + self.assertAlmostEqual(sum(components), total, + delta=1e-12 * abs(total), + msg='the squared-order components do not add up ' + 'to the total this backend computes') + def test_standalone_cpp(self): """test that the scalar C++ standalone exporter is working