From 2f8aeaadbdb67ae7becb6092f574743c53249d07 Mon Sep 17 00:00:00 2001 From: Olivier Mattelaer Date: Wed, 5 Aug 2026 23:17:08 +0200 Subject: [PATCH 01/42] exploit index-permutation symmetry in the color basis and color matrix A color matrix entry is the full contraction of two color structures, so it only depends on the relative labelling of the indices. Any permutation of the external indices mapping the color basis onto itself therefore satisfies C[P(i)][P(j)] = C[i][j], and only one row per orbit has to be computed. ColorBasisSymmetry finds those permutations; for g g > n g the action is transitive, so the whole matrix follows from a single row. The matrix is no longer stored entry by entry either: the distinct values are kept once and the (i,j) grid only holds a compact index into them, so the N^2 dictionaries of ColorFactor objects are gone. The never-read inverted_col_matrix dictionary becomes a property built on demand. In ColorBasis, a recycled color factor only has its indices relabelled, and a relabelled simplified expression is still simplified, so it only needs putting back in canonical form rather than running the full simplification over every term. The equivalence is checked once per canonical representation and the old path is kept where it does not hold. ColorFactor.simplify also looks up similar strings through a dictionary instead of scanning what it has accumulated, which was quadratic in the number of terms. Measured on g g > 6g (5040 color structures, 34300 diagrams): color matrix 194s/2126MB -> 7s/92MB, color basis 167s -> 38s, get_color_amplitudes 4599MB -> 1023MB, peak 8.9GB -> 3.4GB. Generated output is unchanged: 1002 source files identical across standalone, madevent, matchbox and standalone_cpp, and matrix.f for g g > 5g is byte for byte the same. Co-Authored-By: Claude Opus 5 --- madgraph/core/color_algebra.py | 31 +- madgraph/core/color_amp.py | 605 +++++++++++++++++++++--- tests/unit_tests/core/test_color_amp.py | 180 +++++++ 3 files changed, 739 insertions(+), 77 deletions(-) diff --git a/madgraph/core/color_algebra.py b/madgraph/core/color_algebra.py index 8b3ac0db06..75fc891444 100755 --- a/madgraph/core/color_algebra.py +++ b/madgraph/core/color_algebra.py @@ -1128,18 +1128,39 @@ def extend_str(self, new_col_fact): for col_str in new_col_fact: self.append_str(col_str) + @staticmethod + def similarity_key(col_str): + """Hashable key which is equal for two color strings exactly when + ColorString.is_similar says they are, i.e. same Nc power, same + imaginary character and same canonical representation (both the + canonical structure and the index replacement dictionary).""" + + canonical, repl_dict = col_str.to_canonical() + return (col_str.Nc_power, col_str.is_imaginary, canonical, + tuple(sorted(repl_dict.items()))) + def simplify(self): """Returns a new color factor where each color string has been - simplified once and similar strings have been added.""" + simplified once and similar strings have been added. + + Similar strings are looked up through a dictionary rather than by + scanning the strings accumulated so far: the color factors appearing + for high multiplicity processes have thousands of terms, and the linear + scan of append_str made this quadratic. Insertion order is preserved, + so the resulting color factor is identical to the one a scan produces.""" new_col_factor = ColorFactor() + similar = {} # Simplify for col_str in self: res = col_str.simplify() - if res: - new_col_factor.extend_str(res) - else: - new_col_factor.append_str(col_str) + for new_str in (res if res else [col_str]): + key = self.similarity_key(new_str) + try: + similar[key].add(new_str) + except KeyError: + similar[key] = new_str + new_col_factor.append(new_str) # Only returns non zero elements return ColorFactor([col_str for col_str in \ diff --git a/madgraph/core/color_amp.py b/madgraph/core/color_amp.py index 55d5731fef..a849d520b5 100755 --- a/madgraph/core/color_amp.py +++ b/madgraph/core/color_amp.py @@ -21,6 +21,7 @@ import collections import copy import fractions +import itertools import operator import re import array @@ -53,6 +54,13 @@ class ColorBasis(dict): # Dictionary store the raw colorize information _list_color_dict = [] + # Whether relabel_canonical may take its shortcut, per canonical form + _fast_relabel_dict = {} + + # Color objects whose canonical form is fully determined by + # permute_immutable (Tr is cyclic, T is an open chain, ColorOne is empty). + fast_relabel_objects = frozenset(['Tr', 'T', 'ColorOne']) + class ColorBasisError(Exception): """Exception raised if an error occurs in the definition @@ -237,6 +245,66 @@ def add_vertex(self, vertex, diagram, model, return (min_index, new_res_dict) + def _fast_relabel_possible(self, col_fact): + """The shortcut is only attempted on color factors made of objects + whose canonical form is known, and whose indices are all distinct + within a string so that no contraction identity can fire.""" + + for col_str in col_fact: + indices = [] + for name, idx in col_str.to_immutable(): + if name not in self.fast_relabel_objects: + return False + indices.extend(idx) + if len(indices) != len(dict.fromkeys(indices)): + return False + return True + + @staticmethod + def _canonicalize_strings(col_fact): + """Put every color string of col_fact back in canonical form in place + and drop the vanishing ones, mirroring what ColorFactor.simplify does + for an expression which is already simplified.""" + + for col_str in col_fact: + immutable = col_str.to_immutable() + canonical = permute_immutable(immutable, {}) + if canonical != immutable: + col_str.from_immutable(canonical) + col_str.immutable = None + col_str.canonical = None + return color_algebra.ColorFactor([col_str for col_str in col_fact \ + if col_str.coeff != 0]) + + def relabel_canonical(self, col_fact, canonical_rep): + """Return col_fact, which is an already simplified color factor with + relabelled indices, put back in canonical form. Equivalent to + col_fact.simplify().simplify(); which of the two is used is decided + once per canonical representation by running both and comparing.""" + + verdict = self._fast_relabel_dict.get(canonical_rep) + if verdict is True: + return self._canonicalize_strings(col_fact) + if verdict is False: + return col_fact.simplify().simplify() + + # First time this color structure is recycled: check the shortcut + # against the full simplification before trusting it. + slow = col_fact.create_copy().simplify().simplify() + if not self._fast_relabel_possible(col_fact): + self._fast_relabel_dict[canonical_rep] = False + return slow + fast = self._canonicalize_strings(col_fact) + verdict = len(fast) == len(slow) and \ + all(f.to_immutable() == s.to_immutable() and + f.coeff == s.coeff and + f.is_imaginary == s.is_imaginary and + f.Nc_power == s.Nc_power and + f.loop_Nc_power == s.loop_Nc_power + for f, s in zip(fast, slow)) + self._fast_relabel_dict[canonical_rep] = verdict + return fast if verdict else slow + def update_color_basis(self, colorize_dict, index): """Update the current color basis by adding information from the colorize dictionary (produced by the colorize routine) @@ -283,8 +351,18 @@ def update_color_basis(self, colorize_dict, index): # can appear with a loop) to put traces in a canonical ordering. # If it still causes issue, just do a full_simplify(), it would # not bring any heavy additional computational load. - col_fact = col_fact.simplify().simplify() - + # + # What is recycled here is an already simplified color factor + # to which nothing but a relabelling of the indices has been + # applied. A relabelled simplified expression is still + # simplified, so this only has to put every color string back + # in canonical form, which relabel_canonical does directly + # instead of running the full simplification machinery over + # every term. The equivalence of the two is checked once per + # canonical representation, and the slow path is kept for any + # color structure where it does not hold. + col_fact = self.relabel_canonical(col_fact, canonical_rep) + # Here we need to force a specific order for the summed indices # in case we have K6 or K6bar Clebsch Gordan coefficients for colstr in col_fact: colstr.order_summation() @@ -346,6 +424,9 @@ def __init__(self, *args): # Dictionary store the raw colorize information self._list_color_dict = [] + # Whether relabel_canonical may take its shortcut, per canonical form + self._fast_relabel_dict = {} + if args: assert isinstance(args[0], diagram_generation.Amplitude), \ @@ -529,106 +610,479 @@ def color_flow_decomposition(self, repr_dict, ninitial): return res +#=============================================================================== +# Permutation symmetry of a color basis +#=============================================================================== +def permute_immutable(struct, perm): + """Apply the index permutation perm (a dict {old_index: new_index}) to the + immutable representation of a color structure, and bring the result back to + the canonical form used as a ColorBasis key: traces are cyclic, so they are + rotated to start on their smallest index, and the color objects are sorted + exactly as ColorString.to_immutable does.""" + + res = [] + for name, indices in struct: + new_indices = tuple([perm.get(i, i) for i in indices]) + if name == 'Tr' and len(new_indices) > 1: + # Tr is cyclic: rotate so that the smallest index comes first + start = min(range(len(new_indices)), key=new_indices.__getitem__) + new_indices = new_indices[start:] + new_indices[:start] + res.append((name, new_indices)) + res.sort() + return tuple(res) + + +class ColorBasisSymmetry(object): + """Permutations of the external color indices which map a color basis (or a + pair of color bases, for an asymmetric color matrix) onto itself. + + A color matrix entry is the full contraction of two color structures, so it + depends only on the *relative* labelling of the indices: relabelling the + indices consistently in both structures leaves the entry unchanged. Hence + for any such permutation P, + + C[P(i)][P(j)] = C[i][j] + + and only one row per orbit of P-action on the basis has to be computed. + + Note that the permutations found here are not required to be physical + permutations of identical particles: any index relabelling that maps the + basis onto itself is a symmetry of the color matrix. For g g > n g this + finds the full S_(n+2) rather than only the S_n of the final state, which + collapses the whole matrix to a single row.""" + + # Indices above this value are summed indices introduced internally + # (order_summation starts at 10000, colorize uses values below -1000); + # only genuine external indices are permuted. + max_external_index = 1000 + + def __init__(self, keys1, keys2=None): + """keys1/keys2 are the *sorted* lists of color basis keys, i.e. exactly + the ordering used to index the color matrix.""" + + self.keys1 = keys1 + self.keys2 = keys2 if keys2 is not None else keys1 + # permutation of the basis indices induced by each accepted generator + self.generators1 = [] + self.generators2 = [] + # representative of the orbit each row belongs to, and how to get there + # in one step: (parent row, index of the generator mapping it to here) + self.row_rep = list(range(len(keys1))) + self.row_parent = [None] * len(keys1) + self.representatives = list(range(len(keys1))) + + if not keys1 or not self.keys2: + return + + self._find_generators() + self._build_orbits() + + def _external_indices(self, keys): + """Return the sorted list of indices which may be permuted. A plain + list is used rather than a set since 'set' is shadowed by an ordered + variant in this module when reproducible ordering is requested.""" + + indices = {} + for struct in keys: + for _, idx in struct: + for i in idx: + if 0 < i < self.max_external_index: + indices[i] = True + return sorted(indices) + + def _index_signature(self, keys): + """Group indices by the way they appear in the basis: two indices can + only be exchanged if they occupy the same kind of slots. This is only + used to avoid testing hopeless candidates; every candidate is verified + explicitly afterwards.""" + + sig = collections.defaultdict(collections.Counter) + for struct in keys: + for name, idx in struct: + for pos, i in enumerate(idx): + sig[i][(name, len(idx), pos)] += 1 + return sig + + def _find_generators(self): + """Find transpositions of external indices mapping every basis onto + itself, and store the induced permutation of the basis indices.""" + + candidates = self._external_indices(self.keys1) + if self.keys2 is not self.keys1: + other = dict((i, True) for i in self._external_indices(self.keys2)) + candidates = [i for i in candidates if i in other] + if len(candidates) < 2: + return + + sig1 = self._index_signature(self.keys1) + sig2 = self._index_signature(self.keys2) \ + if self.keys2 is not self.keys1 else sig1 + + pos1 = dict((k, i) for i, k in enumerate(self.keys1)) + pos2 = pos1 if self.keys2 is self.keys1 else \ + dict((k, i) for i, k in enumerate(self.keys2)) + + for a, b in itertools.combinations(candidates, 2): + if sig1[a] != sig1[b] or sig2[a] != sig2[b]: + continue + perm = {a: b, b: a} + induced1 = self._induced_permutation(self.keys1, pos1, perm) + if induced1 is None: + continue + if self.keys2 is self.keys1: + induced2 = induced1 + else: + induced2 = self._induced_permutation(self.keys2, pos2, perm) + if induced2 is None: + continue + # A transposition is its own inverse, and so is the permutation it + # induces on the basis. Rows are gathered from their parent with + # the generator itself rather than with its inverse, so make sure + # of it instead of assuming it. + if any(induced1[induced1[i]] != i for i in range(len(induced1))) or \ + any(induced2[induced2[i]] != i for i in range(len(induced2))): + continue + self.generators1.append(induced1) + self.generators2.append(induced2) + + @staticmethod + def _induced_permutation(keys, positions, perm): + """Return the permutation of the basis indices induced by the index + permutation perm, or None if the basis is not mapped onto itself.""" + + induced = [0] * len(keys) + seen = [False] * len(keys) + for i, struct in enumerate(keys): + try: + j = positions[permute_immutable(struct, perm)] + except KeyError: + return None + if seen[j]: + return None + seen[j] = True + induced[i] = j + return induced + + def _build_orbits(self): + """Breadth-first exploration of each orbit, recording for every row the + representative it comes from and the generator that reaches it from its + parent, so that the row can be obtained by a single gather.""" + + if not self.generators1: + return + + n = len(self.keys1) + self.row_rep = [-1] * n + self.representatives = [] + for start in range(n): + if self.row_rep[start] != -1: + continue + self.representatives.append(start) + self.row_rep[start] = start + self.row_parent[start] = None + queue = collections.deque([start]) + while queue: + current = queue.popleft() + for gen_index, induced in enumerate(self.generators1): + image = induced[current] + if self.row_rep[image] == -1: + self.row_rep[image] = start + self.row_parent[image] = (current, gen_index) + queue.append(image) + + def has_symmetry(self): + """True if the symmetry actually reduces the number of rows.""" + + return bool(self.generators1) and \ + len(self.representatives) < len(self.keys1) #=============================================================================== # ColorMatrix #=============================================================================== +class _ColorMatrixView(object): + """Read-only mapping presenting one of the two representations stored by a + ColorMatrix (the ColorFactor one, or the fixed Nc one) as the dictionary + keyed by (i1, i2) that it used to be.""" + + def __init__(self, matrix, entry): + self._matrix = matrix + self._entry = entry + + def __getitem__(self, key): + return self._matrix._get_entry(key)[self._entry] + + def __len__(self): + return len(self._matrix) + + def __contains__(self, key): + try: + self[key] + except (KeyError, IndexError, TypeError): + return False + return True + + def __iter__(self): + return iter(self._matrix) + + def keys(self): + return list(self) + + def values(self): + return [self[key] for key in self] + + def items(self): + return [(key, self[key]) for key in self] + + def get(self, key, default=None): + try: + return self[key] + except (KeyError, IndexError, TypeError): + return default + + def __eq__(self, other): + if isinstance(other, _ColorMatrixView): + if len(self) != len(other): + return False + return all(self[key] == other[key] for key in self) + if isinstance(other, dict): + return dict(self.items()) == other + return NotImplemented + + def __ne__(self, other): + result = self.__eq__(other) + return result if result is NotImplemented else not result + + class ColorMatrix(dict): - """A color matrix, meaning a dictionary with pairs (i,j) as keys where i + """A color matrix, meaning a mapping with pairs (i,j) as keys where i and j refer to elements of color basis objects. Values are Color Factor - objects. Also contains two additional dictionaries, one with the fixed Nc - representation of the matrix, and the other one with the "inverted" matrix, - i.e. a dictionary where keys are values of the color matrix.""" + objects. The fixed Nc representation is available through the + col_matrix_fixed_Nc attribute. + + The matrix is not stored entry by entry. A color matrix entry is the full + contraction of two color structures and therefore only depends on the + relative labelling of the color indices, so entries repeat massively: the + distinct values are stored once and the (i,j) grid only keeps an index into + them. On top of that, index permutations mapping the color basis onto + itself (see ColorBasisSymmetry) relate whole rows to each other, so only + one row per orbit is actually computed; the others are a gather away.""" _col_basis1 = None _col_basis2 = None col_matrix_fixed_Nc = {} - inverted_col_matrix = {} def __init__(self, col_basis, col_basis2=None, Nc=3, Nc_power_min=None, Nc_power_max=None): """Initialize a color matrix with one or two color basis objects. If only one color basis is given, the other one is assumed to be equal. - As options, any value of Nc and minimal/maximal power of Nc can also be + As options, any value of Nc and minimal/maximal power of Nc can also be provided. Note that the min/max power constraint is applied only at the end, so that it does NOT speed up the calculation.""" - self.col_matrix_fixed_Nc = {} - self.inverted_col_matrix = {} - + # Distinct entries, as (result, result_fixed_Nc) pairs, and the (i1,i2) + # grid of indices into that list, stored row-major in a compact array. + self._values = [] + self._val_index = array.array('i') + self._sorted_keys1 = [] + self._sorted_keys2 = [] + self.col_matrix_fixed_Nc = _ColorMatrixView(self, 1) + self._col_basis1 = col_basis if col_basis2: self._col_basis2 = col_basis2 self.build_matrix(Nc, Nc_power_min, Nc_power_max) else: self._col_basis2 = col_basis - # If the two color basis are equal, assumes the color matrix is + # If the two color basis are equal, assumes the color matrix is # symmetric self.build_matrix(Nc, Nc_power_min, Nc_power_max, is_symmetric=True) + #=========================================================================== + # Mapping interface + #=========================================================================== + def _get_entry(self, key): + """Return the (result, result_fixed_Nc) pair for the (i1, i2) key.""" + + i1, i2 = key + n1, n2 = len(self._sorted_keys1), len(self._sorted_keys2) + if not 0 <= i1 < n1 or not 0 <= i2 < n2: + raise KeyError(key) + return self._values[self._val_index[i1 * n2 + i2]] + + def __getitem__(self, key): + return self._get_entry(key)[0] + + def __len__(self): + return len(self._sorted_keys1) * len(self._sorted_keys2) + + def __bool__(self): + return bool(self._sorted_keys1) and bool(self._sorted_keys2) + + __nonzero__ = __bool__ + + def __contains__(self, key): + try: + self._get_entry(key) + except (KeyError, IndexError, TypeError): + return False + return True + + def __iter__(self): + for i1 in range(len(self._sorted_keys1)): + for i2 in range(len(self._sorted_keys2)): + yield (i1, i2) + + def keys(self): + return list(self) + + def values(self): + return [self[key] for key in self] + + def items(self): + return [(key, self[key]) for key in self] + + def get(self, key, default=None): + try: + return self[key] + except (KeyError, IndexError, TypeError): + return default + + def __eq__(self, other): + if isinstance(other, ColorMatrix): + if self._sorted_keys1 != other._sorted_keys1 or \ + self._sorted_keys2 != other._sorted_keys2: + return False + return all(self._get_entry(key) == other._get_entry(key) + for key in self) + if isinstance(other, dict): + return dict(self.items()) == other + return NotImplemented + + def __ne__(self, other): + result = self.__eq__(other) + return result if result is NotImplemented else not result + + __hash__ = None + + @property + def inverted_col_matrix(self): + """Dictionary mapping each fixed Nc value to the list of (i1,i2) it + appears at. Kept for backward compatibility, built on demand.""" + + inverted = {} + for key in self: + inverted.setdefault(self._get_entry(key)[1], []).append(key) + return inverted + + #=========================================================================== + # Construction + #=========================================================================== + def _value_index(self, struct1, struct2, canonical_dict, + Nc, Nc_power_min, Nc_power_max): + """Return the index in self._values of the entry for the two given + color structures, computing it if it is seen for the first time.""" + + # Fix indices in struct2 knowing summed indices in struct1 + # to avoid duplicates + new_struct2 = self.fix_summed_indices(struct1, struct2) + + # Build a canonical representation of the two immutable struct + canonical_entry, dummy = \ + color_algebra.ColorString().to_canonical(struct1 + \ + new_struct2) + + try: + # If this has already been calculated, use the result + return canonical_dict[canonical_entry] + except KeyError: + pass + + # Otherwise calculate the result + result, result_fixed_Nc = self.create_new_entry(struct1, + new_struct2, + Nc_power_min, + Nc_power_max, + Nc) + index = len(self._values) + self._values.append((result, result_fixed_Nc)) + canonical_dict[canonical_entry] = index + return index + def build_matrix(self, Nc=3, Nc_power_min=None, Nc_power_max=None, is_symmetric=False): """Create the matrix using internal color basis objects. Use the stored color basis objects and takes Nc and Nc_min/max parameters as __init__. - If is_isymmetric is True, build only half of the matrix which is assumed - to be symmetric.""" + If is_symmetric is True, the matrix is assumed to be symmetric so that + only half of it needs to be computed.""" - canonical_dict = {} - - for i1, struct1 in \ - enumerate(sorted(self._col_basis1.keys())): - for i2, struct2 in \ - enumerate(sorted(self._col_basis2.keys())): - # Only scan upper right triangle if symmetric - if is_symmetric and i2 < i1: - continue - - # Fix indices in struct2 knowing summed indices in struct1 - # to avoid duplicates - new_struct2 = self.fix_summed_indices(struct1, struct2) + self._sorted_keys1 = sorted(self._col_basis1.keys()) + if self._col_basis2 is self._col_basis1: + self._sorted_keys2 = self._sorted_keys1 + else: + self._sorted_keys2 = sorted(self._col_basis2.keys()) - # Build a canonical representation of the two immutable struct - canonical_entry, dummy = \ - color_algebra.ColorString().to_canonical(struct1 + \ - new_struct2) + keys1, keys2 = self._sorted_keys1, self._sorted_keys2 + n1, n2 = len(keys1), len(keys2) + self._values = [] + self._val_index = array.array('i', [0]) * (n1 * n2) if n1 * n2 else \ + array.array('i') + if not n1 or not n2: + return - try: - # If this has already been calculated, use the result - result, result_fixed_Nc = canonical_dict[canonical_entry] - except KeyError: - # Otherwise calculate the result - result, result_fixed_Nc = \ - self.create_new_entry(struct1, - new_struct2, - Nc_power_min, - Nc_power_max, - Nc) - # Store both results - canonical_dict[canonical_entry] = (result, result_fixed_Nc) - - # Store the full result... - self[(i1, i2)] = result - if is_symmetric: - self[(i2, i1)] = result - - # the fixed Nc one ... - self.col_matrix_fixed_Nc[(i1, i2)] = result_fixed_Nc - if is_symmetric: - self.col_matrix_fixed_Nc[(i2, i1)] = result_fixed_Nc - # and update the inverted dict - if result_fixed_Nc in list(self.inverted_col_matrix.keys()): - self.inverted_col_matrix[result_fixed_Nc].append((i1, - i2)) - if is_symmetric: - self.inverted_col_matrix[result_fixed_Nc].append((i2, - i1)) - else: - self.inverted_col_matrix[result_fixed_Nc] = [(i1, i2)] + canonical_dict = {} + symmetry = ColorBasisSymmetry(keys1, + None if keys2 is keys1 else keys2) + + if not symmetry.has_symmetry(): + # No index permutation maps the basis onto itself: fall back to the + # plain scan, using the symmetry of the matrix itself if available. + for i1, struct1 in enumerate(keys1): + for i2, struct2 in enumerate(keys2): + if is_symmetric and i2 < i1: + continue + index = self._value_index(struct1, struct2, canonical_dict, + Nc, Nc_power_min, Nc_power_max) + self._val_index[i1 * n2 + i2] = index if is_symmetric: - self.inverted_col_matrix[result_fixed_Nc] = [(i2, i1)] + self._val_index[i2 * n2 + i1] = index + return + + # One row per orbit is computed explicitly; every other row is the + # image of an already known one under a single generator. + done = [False] * n1 + for rep in symmetry.representatives: + struct1 = keys1[rep] + offset = rep * n2 + for i2, struct2 in enumerate(keys2): + self._val_index[offset + i2] = \ + self._value_index(struct1, struct2, canonical_dict, + Nc, Nc_power_min, Nc_power_max) + done[rep] = True + + # Breadth-first replay of the orbit exploration: a row whose parent is + # already filled is obtained by permuting the parent's columns. + remaining = [i for i in range(n1) if not done[i]] + while remaining: + progressed = False + still_missing = [] + for row in remaining: + parent, gen_index = symmetry.row_parent[row] + if not done[parent]: + still_missing.append(row) + continue + induced2 = symmetry.generators2[gen_index] + src = parent * n2 + dest = row * n2 + val_index = self._val_index + for i2 in range(n2): + val_index[dest + i2] = val_index[src + induced2[i2]] + done[row] = True + progressed = True + assert progressed, "Color matrix orbit exploration made no progress" + remaining = still_missing def create_new_entry(self, struct1, struct2, Nc_power_min, Nc_power_max, Nc): @@ -690,25 +1144,32 @@ def __str__(self): return mystr + def _fixed_Nc_row(self, line_index): + """Return the fixed Nc entries of one line of the matrix.""" + + n2 = len(self._sorted_keys2) + offset = line_index * n2 + values = self._values + val_index = self._val_index + return [values[val_index[offset + i2]][1] for i2 in range(n2)] + def get_line_denominators(self): """Get a list with the denominators for the different lines in the color matrix""" den_list = [] - for i1 in range(len(self._col_basis1)): - den_list.append(self.lcmm(*[\ - self.col_matrix_fixed_Nc[(i1, i2)][0].denominator for \ - i2 in range(len(self._col_basis2))])) - + for i1 in range(len(self._sorted_keys1)): + den_list.append(self.lcmm(*[entry[0].denominator for entry in \ + self._fixed_Nc_row(i1)])) + return den_list def get_line_numerators(self, line_index, den): """Returns a list of numerator for line line_index, assuming a common denominator den.""" - return [self.col_matrix_fixed_Nc[(line_index, i2)][0].numerator * \ - den / self.col_matrix_fixed_Nc[(line_index, i2)][0].denominator \ - for i2 in range(len(self._col_basis2))] + return [entry[0].numerator * den / entry[0].denominator \ + for entry in self._fixed_Nc_row(line_index)] @classmethod def fix_summed_indices(self, struct1, struct2): diff --git a/tests/unit_tests/core/test_color_amp.py b/tests/unit_tests/core/test_color_amp.py index 9775c450fc..c0dd0c7719 100755 --- a/tests/unit_tests/core/test_color_amp.py +++ b/tests/unit_tests/core/test_color_amp.py @@ -765,3 +765,183 @@ def test_helper_lcm_functions(self): self.assertEqual(color_amp.ColorMatrix.lcm(6, 3), 6) self.assertEqual(color_amp.ColorMatrix.lcmm(6, 3, 5, 2), 30) + + def get_gluon_amplitude(self, n_final): + """Amplitude for g g > n_final gluons, using the test model.""" + + myleglist = base_objects.LegList() + myleglist.append(base_objects.Leg({'id':21, 'state':False})) + myleglist.append(base_objects.Leg({'id':21, 'state':False})) + myleglist.extend([base_objects.Leg({'id':21, + 'state':True})] * n_final) + myamplitude = diagram_generation.Amplitude() + myamplitude.set('process', base_objects.Process({'legs':myleglist, + 'model':self.mymodel})) + myamplitude.generate_diagrams() + return myamplitude + + def get_quark_amplitude(self, n_pairs): + """Amplitude for u u~ > n_pairs (u u~), using the test model.""" + + myleglist = base_objects.LegList() + myleglist.append(base_objects.Leg({'id':2, 'state':False})) + myleglist.append(base_objects.Leg({'id':-2, 'state':False})) + for _ in range(n_pairs): + myleglist.append(base_objects.Leg({'id':2, 'state':True})) + myleglist.append(base_objects.Leg({'id':-2, 'state':True})) + myamplitude = diagram_generation.Amplitude() + myamplitude.set('process', base_objects.Process({'legs':myleglist, + 'model':self.mymodel})) + myamplitude.generate_diagrams() + return myamplitude + + def test_permute_immutable(self): + """Test the canonical form of a permuted color structure: traces are + cyclic so they are rotated to start on their smallest index, open + chains are not, and the color objects are sorted.""" + + # Tr is cyclic: exchanging 1 and 3 has to be rotated back + self.assertEqual(color_amp.permute_immutable((('Tr', (1, 2, 3, 4)),), + {1: 3, 3: 1}), + (('Tr', (1, 4, 3, 2)),)) + # identity permutation still canonicalises + self.assertEqual(color_amp.permute_immutable((('Tr', (3, 4, 1, 2)),), + {}), + (('Tr', (1, 2, 3, 4)),)) + # T is an open chain, no rotation, but the factors get sorted + self.assertEqual(color_amp.permute_immutable((('T', (5, 2, 1)), + ('T', (5, 4, 3))), + {1: 3, 3: 1}), + (('T', (5, 2, 3)), ('T', (5, 4, 1)))) + # indices which are not in the permutation are left alone + self.assertEqual(color_amp.permute_immutable((('ColorOne', ()),), {}), + (('ColorOne', ()),)) + + def test_color_basis_symmetry(self): + """The permutations found have to map the basis onto itself, and for + a pure gluon process they act transitively so a single row of the + color matrix determines all the others.""" + + for n_final in range(1, 4): + col_basis = color_amp.ColorBasis(self.get_gluon_amplitude(n_final)) + keys = sorted(col_basis.keys()) + symmetry = color_amp.ColorBasisSymmetry(keys) + + positions = dict((key, i) for i, key in enumerate(keys)) + for induced in symmetry.generators1: + # a permutation of the basis, and an involution + self.assertEqual(sorted(induced), list(range(len(keys)))) + for i in range(len(keys)): + self.assertEqual(induced[induced[i]], i) + + # all gluons are equivalent, so there is a single orbit + self.assertEqual(len(symmetry.representatives), 1) + # every row is reachable from the representative + self.assertEqual(symmetry.row_rep, [0] * len(keys)) + self.assertEqual(len(positions), len(keys)) + + def test_color_basis_symmetry_no_symmetry(self): + """When no index permutation maps the basis onto itself, every row is + its own representative and the color matrix falls back to the plain + scan, which still has to give a correct symmetric matrix.""" + + # no transposition maps these bases onto themselves + for keys in [[(('T', (1, 2, 3)),), (('Tr', (1, 2, 3)),)], + [(('T', (1, 2, 3)),)]]: + symmetry = color_amp.ColorBasisSymmetry(keys) + self.assertEqual(symmetry.generators1, []) + self.assertFalse(symmetry.has_symmetry()) + self.assertEqual(symmetry.representatives, list(range(len(keys)))) + + # the color matrix then falls back to the plain scan + keys = [(('T', (1, 2, 3)),)] + col_matrix = color_amp.ColorMatrix(dict((key, []) for key in keys), + Nc=3) + self.assertEqual(len(col_matrix), 1) + self.assertEqual(col_matrix.col_matrix_fixed_Nc[(0, 0)], + (fractions.Fraction(4), 0)) + + def test_color_matrix_matches_direct_computation(self): + """The color matrix built by orbits must agree entry by entry with the + one obtained by computing every entry independently.""" + + for amplitude in [self.get_gluon_amplitude(2), + self.get_gluon_amplitude(3), + self.get_quark_amplitude(1)]: + col_basis = color_amp.ColorBasis(amplitude) + col_matrix = color_amp.ColorMatrix(col_basis, Nc=3) + keys = sorted(col_basis.keys()) + + for i1, struct1 in enumerate(keys): + for i2, struct2 in enumerate(keys): + new_struct2 = color_amp.ColorMatrix.fix_summed_indices( + struct1, struct2) + result, result_fixed_Nc = col_matrix.create_new_entry( + struct1, new_struct2, None, None, 3) + self.assertEqual(col_matrix.col_matrix_fixed_Nc[(i1, i2)], + result_fixed_Nc) + # entries are recycled between equivalent index pairs, so + # the terms of the color factor can come in another order + self.assertEqual(sorted(map(str, col_matrix[(i1, i2)])), + sorted(map(str, result))) + + # the mapping interface still behaves like the dictionary it was + self.assertEqual(len(col_matrix), len(keys) ** 2) + self.assertTrue(col_matrix) + self.assertTrue((0, 0) in col_matrix) + self.assertFalse((0, len(keys)) in col_matrix) + self.assertEqual(sorted(col_matrix.keys()), + sorted([(i, j) for i in range(len(keys)) + for j in range(len(keys))])) + + def test_relabel_canonical(self): + """The shortcut taken when a simplified color factor is recycled with + relabelled indices must agree with the full simplification, and it has + to be the path actually taken for QCD processes.""" + + col_basis = color_amp.ColorBasis() + col_basis.build(self.get_gluon_amplitude(3)) + self.assertTrue(col_basis._fast_relabel_dict) + self.assertTrue(all(col_basis._fast_relabel_dict.values())) + + # explicitly compare both paths on every recycled structure + for color_dict in col_basis._list_color_dict: + for col_str in color_dict.values(): + canonical_rep, rep_dict = col_str.to_canonical() + if canonical_rep not in col_basis._canonical_dict: + continue + col_fact = col_basis._canonical_dict[canonical_rep].create_copy() + col_fact.replace_indices(col_basis._invert_dict(rep_dict)) + for one_str in col_fact: + one_str.coeff = one_str.coeff * col_str.coeff + slow = col_fact.create_copy().simplify().simplify() + fast = col_basis._canonicalize_strings(col_fact) + self.assertEqual([s.to_immutable() for s in fast], + [s.to_immutable() for s in slow]) + self.assertEqual([s.coeff for s in fast], + [s.coeff for s in slow]) + + def test_color_factor_simplify_merges_like_strings(self): + """ColorFactor.simplify has to add up similar strings, keeping them in + order of first appearance.""" + + col_fact = color.ColorFactor([ + color.ColorString([color.T(1, 2, 3)], + coeff=fractions.Fraction(2, 3)), + color.ColorString([color.T(4, 5, 6)], + coeff=fractions.Fraction(1, 5)), + color.ColorString([color.T(1, 2, 3)], + coeff=fractions.Fraction(1, 3))]) + result = col_fact.simplify() + self.assertEqual([col_str.to_immutable() for col_str in result], + [(('T', (1, 2, 3)),), (('T', (4, 5, 6)),)]) + self.assertEqual([col_str.coeff for col_str in result], + [fractions.Fraction(1, 1), fractions.Fraction(1, 5)]) + + # strings adding up to zero are dropped + col_fact = color.ColorFactor([ + color.ColorString([color.T(1, 2, 3)], + coeff=fractions.Fraction(1, 3)), + color.ColorString([color.T(1, 2, 3)], + coeff=fractions.Fraction(-1, 3))]) + self.assertEqual(len(col_fact.simplify()), 0) From dab4a8db7f621730a25590acff5351ce2827c63c Mon Sep 17 00:00:00 2001 From: Olivier Mattelaer Date: Thu, 6 Aug 2026 07:28:18 +0200 Subject: [PATCH 02/42] index the JAMP matrix instead of scanning every column of it optimise_jamp looks for pairs of columns whose ratio is the same on many lines. It found the candidates by looking every column of the 0..nb_col range up in the matrix, for every non zero entry. That matrix is very sparse - a color flow only gets a small share of the amplitudes, 322 out of 7245 for g g > 5g - so nearly all of that work was spent discovering zeros, and the cost was the number of non zero entries times the number of columns. Index the non zero entries by line instead, and walk only those. The pairs are visited in the same order as before, so the sub-expressions found and the order they are defined in are unchanged. The substitution step gets the same treatment through an index by column, and the count of lines sharing a ratio now lives in a single dictionary keyed by the two columns and the ratio at once rather than in nested dictionaries. g g > 5g: get_JAMP_lines 33.5s -> 5.2s, and the whole standalone output 29.3s -> 14.3s. g g > 6g (5040 color flows, 126630 amplitudes, 8.1M non zero entries) now goes through in 270s where the scan was out of reach. Generated output is unchanged: the same 1002 source files across standalone, madevent, matchbox and standalone_cpp, and matrix.f for g g > 5g is byte for byte the same as before. Co-Authored-By: Claude Opus 5 --- madgraph/iolibs/export_v4.py | 83 ++++++++++-- tests/unit_tests/iolibs/test_export_v4.py | 154 +++++++++++++++++++++- 2 files changed, 223 insertions(+), 14 deletions(-) diff --git a/madgraph/iolibs/export_v4.py b/madgraph/iolibs/export_v4.py index 525478ac5c..5e2823e4c6 100755 --- a/madgraph/iolibs/export_v4.py +++ b/madgraph/iolibs/export_v4.py @@ -17,6 +17,7 @@ from madgraph.core import base_objects """Methods and classes to export matrix elements to v4 format.""" +import bisect import copy import math, cmath from io import StringIO @@ -2548,6 +2549,45 @@ def format(frac): return res_list, len(defs) + @staticmethod + def index_jamp_matrix(all_element, nb_col): + """Sorted lists of the positions of the non zero entries of the matrix, + by line and by column. An entry which is present but zero does not + count, and neither does a column outside the 0..nb_col range, so that + these indices list exactly the entries the plain scan would look at.""" + + lines = collections.defaultdict(list) + columns = collections.defaultdict(list) + for (i, j), value in all_element.items(): + if value and j < nb_col: + lines[i].append(j) + columns[j].append(i) + for line in lines.values(): + line.sort() + for column in columns.values(): + column.sort() + return lines, columns + + @staticmethod + def common_jamp_lines(columns, nb_line, j1, j2): + """Lines, in increasing order, where both columns j1 and j2 are non + zero. Both column lists are sorted, so this is a plain merge.""" + + left, right = columns.get(j1, []), columns.get(j2, []) + res = [] + pos1 = pos2 = 0 + while pos1 < len(left) and pos2 < len(right): + if left[pos1] == right[pos2]: + if left[pos1] < nb_line: + res.append(left[pos1]) + pos1 += 1 + pos2 += 1 + elif left[pos1] < right[pos2]: + pos1 += 1 + else: + pos2 += 1 + return res + def optimise_jamp(self, all_element, nb_line=0, nb_col=0, added=0): """ optimise problem of type Y = A X A is a matrix (all_element) @@ -2597,18 +2637,32 @@ def optimise_jamp(self, all_element, nb_line=0, nb_col=0, added=0): newdef1 = newdef1 + new_def return all_element, newdef1 + # Index of the non zero entries, by line and by column. The matrix is + # very sparse (a color flow only gets a small share of the amplitudes) + # so walking the whole 0..nb_col range for every entry, as looking the + # columns up one by one in the matrix amounts to, spends nearly all of + # its time discovering zeros. + lines, columns = self.index_jamp_matrix(all_element, nb_col) + max_count = 0 all_index = [] - operation = collections.defaultdict(lambda: collections.defaultdict(int)) + # how many lines have the same ratio between two given columns, keyed + # by the two columns and the ratio at once rather than by nested + # dictionaries: this is the innermost loop of the whole optimisation + operation = collections.defaultdict(int) for (i,j1), v1 in all_element.items(): - ratios = [(j2,all_element.get((i,j2), 0)/v1) for j2 in range(j1+1, nb_col) if all_element.get((i,j2), 0)] - for j2, R in ratios: - operation[(j1,j2)][R] +=1 - if operation[(j1,j2)][R] > max_count: - max_count = operation[(j1,j2)][R] - all_index = [(j1,j2, R)] - elif operation[(j1,j2)][R] == max_count: - all_index.append((j1,j2, R)) + line = lines.get(i) + if not line: + continue + for j2 in line[bisect.bisect_right(line, j1):]: + key = (j1, j2, all_element[(i,j2)]/v1) + operation[key] += 1 + count = operation[key] + if count > max_count: + max_count = count + all_index = [key] + elif count == max_count: + all_index.append(key) if max_count <= 1: return all_element, [] @@ -2617,20 +2671,23 @@ def optimise_jamp(self, all_element, nb_line=0, nb_col=0, added=0): for index in all_index: j1,j2,R = index first = True - for i in range(nb_line): + # only the lines where both columns are filled can contribute; the + # substitutions done here can empty some of them, so the values + # still have to be read back from the matrix + for i in self.common_jamp_lines(columns, nb_line, j1, j2): v1 = all_element.get((i,j1), 0) v2 = all_element.get((i,j2), 0) - if not v1 or not v2: + if not v1 or not v2: continue if v2/v1 == R: if first: first = False added +=1 to_add.append((added,j1,j2,R, max_count)) - + all_element[(i,-added)] = v1 del all_element[(i,j1)] #= 0 - del all_element[(i,j2)] #= 0 + del all_element[(i,j2)] #= 0 logger.log(5,"Define %d new shortcut reused %d times", len(to_add), max_count) new_element, new_def = self.optimise_jamp(all_element, nb_line=nb_line, nb_col=nb_col, added=added) diff --git a/tests/unit_tests/iolibs/test_export_v4.py b/tests/unit_tests/iolibs/test_export_v4.py index f2d500e560..aae1514b2d 100644 --- a/tests/unit_tests/iolibs/test_export_v4.py +++ b/tests/unit_tests/iolibs/test_export_v4.py @@ -16,9 +16,11 @@ """Unit test library for the export v4 format routines""" from __future__ import absolute_import +import collections import copy import fractions -import os +import os +import random import sys root_path = os.path.split(os.path.dirname(os.path.realpath( __file__ )))[0] sys.path.append(os.path.join(root_path, os.path.pardir, os.path.pardir)) @@ -10360,3 +10362,153 @@ def test_madevent_template_uses_decay_aware_broken_symmetry_metadata(self): [me.get('diagrams')[323], me.get('diagrams')[954], me.get('diagrams')[1123], me.get('diagrams')[1139]]) + + +class OptimiseJampTest(unittest.TestCase): + """Test the common sub-expression elimination applied to the JAMP + definitions.""" + + @staticmethod + def reference_optimise_jamp(all_element, nb_line=0, nb_col=0, added=0): + """Straightforward version of ProcessExporterFortran.optimise_jamp, + looking every column up in the matrix instead of indexing the non zero + entries. The splitting of wide matrices is left out, so keep the test + matrices below the 600 columns which trigger it.""" + + if not nb_line: + for i, j in all_element: + if i + 1 > nb_line: + nb_line = i + 1 + if j + 1 > nb_col: + nb_col = j + 1 + assert nb_col <= 600 + + max_count = 0 + all_index = [] + operation = collections.defaultdict( + lambda: collections.defaultdict(int)) + for (i, j1), v1 in all_element.items(): + ratios = [(j2, all_element.get((i, j2), 0) / v1) + for j2 in range(j1 + 1, nb_col) + if all_element.get((i, j2), 0)] + for j2, R in ratios: + operation[(j1, j2)][R] += 1 + if operation[(j1, j2)][R] > max_count: + max_count = operation[(j1, j2)][R] + all_index = [(j1, j2, R)] + elif operation[(j1, j2)][R] == max_count: + all_index.append((j1, j2, R)) + + if max_count <= 1: + return all_element, [] + + to_add = [] + for j1, j2, R in all_index: + first = True + for i in range(nb_line): + v1 = all_element.get((i, j1), 0) + v2 = all_element.get((i, j2), 0) + if not v1 or not v2: + continue + if v2 / v1 == R: + if first: + first = False + added += 1 + to_add.append((added, j1, j2, R, max_count)) + all_element[(i, -added)] = v1 + del all_element[(i, j1)] + del all_element[(i, j2)] + + new_element, new_def = OptimiseJampTest.reference_optimise_jamp( + all_element, nb_line, nb_col, added) + for one_def in to_add: + new_def.insert(0, one_def) + return new_element, new_def + + @staticmethod + def random_matrix(seed, nb_line, nb_col, density): + """Sparse matrix with repeating values, so that the optimisation has + something to find.""" + + # no zero value: the scan divides by the entry it starts from, so a + # stored zero in a line which has other entries makes it raise + values = [1, -1, 2, -2, 0.5, -0.5, 3, 1j, -1j] + rng = random.Random(seed) + all_element = {} + for i in range(nb_line): + for j in range(nb_col): + if rng.random() < density: + all_element[(i, j)] = complex(rng.choice(values)) + return all_element + + def test_index_jamp_matrix(self): + """The indices must list the non zero entries in increasing order, and + leave out both the zero values and the columns beyond nb_col.""" + + all_element = {(0, 2): 1, (0, 0): 3, (0, 1): 0, + (1, 0): 2, (1, 5): 7, (0, -1): 4} + lines, columns = export_v4.ProcessExporterFortran.index_jamp_matrix( + all_element, 3) + self.assertEqual(dict(lines), {0: [-1, 0, 2], 1: [0]}) + self.assertEqual(dict(columns), {-1: [0], 0: [0, 1], 2: [0]}) + + def test_common_jamp_lines(self): + """Lines where two columns are both filled, in increasing order.""" + + columns = {1: [0, 2, 3, 7], 2: [2, 3, 5], 3: [9]} + common = export_v4.ProcessExporterFortran.common_jamp_lines + self.assertEqual(common(columns, 10, 1, 2), [2, 3]) + self.assertEqual(common(columns, 3, 1, 2), [2]) + self.assertEqual(common(columns, 10, 1, 3), []) + self.assertEqual(common(columns, 10, 1, 4), []) + + def test_optimise_jamp_matches_reference(self): + """The optimisation has to give exactly the same sub-expressions as + the straightforward scan, including the order they are defined in.""" + + exporter = export_v4.ProcessExporterFortranSA() + for seed, nb_line, nb_col, density in [(1, 8, 12, 0.5), + (2, 12, 25, 0.35), + (3, 20, 40, 0.25), + (4, 5, 5, 0.9), + (5, 30, 15, 0.6)]: + all_element = self.random_matrix(seed, nb_line, nb_col, density) + reference, reference_def = self.reference_optimise_jamp( + dict(all_element)) + exporter.myjamp_count = 0 + result, result_def = exporter.optimise_jamp(dict(all_element)) + self.assertEqual(result_def, reference_def) + self.assertEqual(result, reference) + + def test_optimise_jamp_stored_zero(self): + """An entry which is present but zero must be ignored, like the scan + looking the columns up in the matrix does.""" + + # the zero has to be the last entry of its line: the scan divides by + # the entry it starts from, so a zero followed by a non zero raises + all_element = {} + for i in range(4): + all_element[(i, 0)] = complex(1) + all_element[(i, 1)] = complex(2) + all_element[(i, 2)] = complex(0) + reference, reference_def = self.reference_optimise_jamp( + dict(all_element)) + exporter = export_v4.ProcessExporterFortranSA() + exporter.myjamp_count = 0 + result, result_def = exporter.optimise_jamp(dict(all_element)) + self.assertEqual(result_def, reference_def) + self.assertEqual(result, reference) + # column 2 is never picked up as a sub-expression + self.assertTrue(result_def) + self.assertTrue(all(2 not in (j1, j2) + for _, j1, j2, _, _ in result_def)) + + def test_optimise_jamp_no_saving(self): + """A matrix where no sub-expression is reused is returned as is.""" + + exporter = export_v4.ProcessExporterFortranSA() + exporter.myjamp_count = 0 + all_element = {(0, 0): complex(1), (1, 1): complex(2)} + result, defs = exporter.optimise_jamp(dict(all_element)) + self.assertEqual(defs, []) + self.assertEqual(result, all_element) From 0911a2d59c9180c66d2c49c77a080f1f313d7319 Mon Sep 17 00:00:00 2001 From: Olivier Mattelaer Date: Thu, 6 Aug 2026 07:42:47 +0200 Subject: [PATCH 03/42] build the powers of Nc in the JAMP coefficients once per power Every amplitude coefficient carries a power of the number of colors, and Fraction(3)**power was built again for each of them: 231840 times for g g > 5g, 8.1M times for g g > 6g, for a handful of distinct powers. This is a small win only. The Fraction time left in get_JAMP_lines is not in building those powers but in the multiplications mixing complex, int and Fraction, which fall back on the reflected operators. Reordering the product to keep it rational until the end would avoid that, but it would also move where the rounding to floating point happens, and optimise_jamp compares the resulting ratios for equality, so the generated code would change. g g > 5g get_JAMP_lines: 3.50s -> 3.46s, within the run to run spread. Generated output unchanged, the same 1002 source files as before. Co-Authored-By: Claude Opus 5 --- madgraph/iolibs/export_v4.py | 11 ++++++++++- 1 file changed, 10 insertions(+), 1 deletion(-) diff --git a/madgraph/iolibs/export_v4.py b/madgraph/iolibs/export_v4.py index 5e2823e4c6..9869c63aee 100755 --- a/madgraph/iolibs/export_v4.py +++ b/madgraph/iolibs/export_v4.py @@ -2409,6 +2409,10 @@ def get_JAMP_lines(self, col_amps, JAMP_format="JAMP(%s)", AMP_format="AMP(%s)", all_element = {} res_list = [] + # Every single amplitude carries a power of the number of colors in its + # coefficient, but a process only uses a handful of distinct powers, so + # build the corresponding fractions once instead of once per amplitude. + nc_powers = {} for i, coeff_list in enumerate(color_amplitudes): # It might happen that coeff_list is empty if this function was # called from get_JAMP_lines_split_order (i.e. if some color flow @@ -2442,7 +2446,12 @@ def get_JAMP_lines(self, col_amps, JAMP_format="JAMP(%s)", AMP_format="AMP(%s)", for (coefficient, amp_number) in coefs: if not coefficient: continue - value = (1j if coefficient[2] else 1)* coefficient[0] * coefficient[1] * fractions.Fraction(3)**coefficient[3] + try: + nc_power = nc_powers[coefficient[3]] + except KeyError: + nc_power = fractions.Fraction(3)**coefficient[3] + nc_powers[coefficient[3]] = nc_power + value = (1j if coefficient[2] else 1)* coefficient[0] * coefficient[1] * nc_power if (i+1, amp_number) not in all_element: all_element[(i+1, amp_number)] = value else: From 390c493eee49c39d7fbadf4e029bd1005b2b7f69 Mon Sep 17 00:00:00 2001 From: Olivier Mattelaer Date: Thu, 6 Aug 2026 09:44:49 +0200 Subject: [PATCH 04/42] encode the color matrix by one line per orbit, and split the contraction sum Two independent changes to the standalone output, both in the color matrix. Every line of the color matrix is the line of its orbit representative with the columns permuted, in every process looked at: pure gluon, one quark line, several quark lines, quark initiated. So instead of writing the N*(N+1)/2 entries out, write one line per orbit, the permutations of a spanning subset of the generators, and for each line which generator reaches it from which other line. INIT_CF follows those back to the representative on the first call and rebuilds the matrix. This is only taken when it is at least four times smaller, which leaves everything below about a hundred color structures on the path it was on. g g > 5g: 259560 -> 5760 numbers, matrix.f 3.5MB -> 2.3MB. g g > 6g: 12703320 -> 45360 numbers, the color block of matrix.f 73MB -> 0.2MB and the file itself 114MB -> 41MB. The contraction of the matrix with the JAMPs now sums over four accumulators rather than one. A single accumulator makes every term wait for the one before it to leave the adder, and that latency, not the arithmetic, is what the loop was spending its time on: 1.55x faster on the generated g g > 5g code, 1.6x to 2x in a standalone benchmark up to 5040 colors. Narrowing the coefficients to one byte or widening them to a double changes nothing, and no compiler flag does this by itself, -ffast-math included. Reordering the sum moves the rounding, so |M|^2 changes in the last bit: checked against the previous output on g g > 4g, g g > tt~ gg, uu~ > uu~ gg, gg > uu~ dd~ g, g g > 3g and g g > 5g, all agreeing to one ulp or exactly. The color basis and color matrix themselves are unchanged. Co-Authored-By: Claude Opus 5 --- madgraph/core/color_amp.py | 57 +++++++ madgraph/iolibs/export_v4.py | 154 ++++++++++++++++++ .../template_files/matrix_standalone_v4.inc | 33 +++- .../matrix.f | 35 +++- tests/unit_tests/core/test_color_amp.py | 52 ++++++ 5 files changed, 319 insertions(+), 12 deletions(-) diff --git a/madgraph/core/color_amp.py b/madgraph/core/color_amp.py index a849d520b5..9250f1a69f 100755 --- a/madgraph/core/color_amp.py +++ b/madgraph/core/color_amp.py @@ -796,6 +796,63 @@ def has_symmetry(self): return bool(self.generators1) and \ len(self.representatives) < len(self.keys1) + def spanning_generators(self): + """Indices of a subset of the generators which still reaches every + line of every orbit. Anything which writes the permutations out has to + store one array of basis indices per generator, so dropping those which + connect nothing new makes a large difference.""" + + parent = list(range(len(self.keys1))) + + def find(x): + while parent[x] != x: + parent[x] = parent[parent[x]] + x = parent[x] + return x + + kept = [] + for index, induced in enumerate(self.generators1): + used = False + for i, j in enumerate(induced): + ri, rj = find(i), find(j) + if ri != rj: + parent[ri] = rj + used = True + if used: + kept.append(index) + return kept + + def spanning_tree(self, gen_indices=None): + """Describe every line as one generator applied to another line: + returns the orbit representatives, the representative of each line, the + (parent line, position in gen_indices) pair reaching each line, and the + generators actually used. Following the parents back to the + representative gives the permutation relating the two lines.""" + + if gen_indices is None: + gen_indices = self.spanning_generators() + gens = [self.generators1[i] for i in gen_indices] + + n = len(self.keys1) + representative = [-1] * n + parent = [None] * n + representatives = [] + for start in range(n): + if representative[start] != -1: + continue + representatives.append(start) + representative[start] = start + queue = collections.deque([start]) + while queue: + current = queue.popleft() + for local, induced in enumerate(gens): + image = induced[current] + if representative[image] == -1: + representative[image] = start + parent[image] = (current, local) + queue.append(image) + return representatives, representative, parent, gens + #=============================================================================== # ColorMatrix diff --git a/madgraph/iolibs/export_v4.py b/madgraph/iolibs/export_v4.py index 9869c63aee..777cd3427d 100755 --- a/madgraph/iolibs/export_v4.py +++ b/madgraph/iolibs/export_v4.py @@ -42,6 +42,7 @@ import models import madgraph.core.base_objects as base_objects import madgraph.core.color_algebra as color +import madgraph.core.color_amp as color_amp import madgraph.core.helas_objects as helas_objects import madgraph.iolibs.drawing_eps as draw import madgraph.iolibs.files as files @@ -226,6 +227,9 @@ class ProcessExporterFortran(VirtualExporter): } grouped_mode = False jamp_optim = False + # how much smaller the compressed color matrix has to be before it is used + # instead of writing every entry out (see get_color_matrix_encoding) + color_encoding_margin = 4 run_card_class = None use_flavor_mask = True @@ -2062,6 +2066,58 @@ def format_integer_list(self, list, name, n=5): + def get_color_matrix_encoding(self, matrix_element): + """Describe the color matrix by one line per orbit of the index + permutations leaving the color basis invariant, plus the permutations + needed to reach every other line from it (see ColorBasisSymmetry). + + Every line of the matrix is one of those lines with its columns + permuted, so this replaces the N*(N+1)/2 entries by (nrep+ngen+3)*N + numbers. That is only a gain once the basis is large enough, and None + is returned otherwise so that the entries are written out as before.""" + + color_matrix = matrix_element.get('color_matrix') + if not color_matrix: + return None + # an asymmetric matrix does not have the line structure exploited here + if color_matrix._col_basis1 is not color_matrix._col_basis2: + return None + + keys = color_matrix._sorted_keys1 + nb_color = len(keys) + symmetry = color_amp.ColorBasisSymmetry(keys) + if not symmetry.has_symmetry(): + return None + representatives, representative, parent, gens = symmetry.spanning_tree() + + # Writing the entries out is well trodden and the compressed form + # carries a routine of its own, so only take it when it pays clearly. + # In practice this leaves everything below about a hundred color + # structures alone, which is where the matrix is not the bulk of the + # generated file anyway. + size = (len(representatives) + len(gens) + 3) * nb_color + if size * self.color_encoding_margin > nb_color * (nb_color + 1) // 2: + return None + + denominator = max(color_matrix.get_line_denominators()) + slot = dict((line, index) for index, line in enumerate(representatives)) + rows = [] + for line in representatives: + num_list = color_matrix.get_line_numerators(line, denominator) + assert all(int(i) == i for i in num_list) + rows.append([int(i) for i in num_list]) + + return {'denom': denominator, + 'nb_color': nb_color, + 'rows': rows, + 'gens': gens, + # for each line, the line it comes from and the generator + # reaching it, or (0,0) when the line is a representative + 'parent': [(0, 0) if p is None else (p[0] + 1, p[1] + 1) + for p in parent], + 'slot': [slot[representative[i]] + 1 + for i in range(nb_color)]} + def get_color_data_lines(self, matrix_element, n=128): """Return the color matrix definition lines for this matrix element. Split rows in chunks of size n.""" @@ -2069,6 +2125,14 @@ def get_color_data_lines(self, matrix_element, n=128): if not matrix_element.get('color_matrix'): return ["DATA %(proc_prefix)sDenom/1/", "DATA %(proc_prefix)sCF/1/"] + if self.get_color_matrix_encoding(matrix_element): + # the entries are rebuilt at run time by INIT_CF, only the overall + # denominator is still needed here + denominator = max(matrix_element.get('color_matrix').\ + get_line_denominators()) + return ["DATA %%(proc_prefix)sDenom/%(denom)i/" % \ + {'denom': denominator}] + ret_list = [] my_cs = color.ColorString() denominator = max(matrix_element.get('color_matrix').get_line_denominators()) @@ -2101,6 +2165,93 @@ def get_color_data_lines(self, matrix_element, n=128): return ret_list + @staticmethod + def get_int_data_lines(name, values, n=128): + """DATA statements filling the one dimensional integer array name.""" + + lines = [] + for start in range(0, len(values), n): + chunk = values[start:start + n] + lines.append(" DATA (%s(i),i=%d,%d) /%s/" % \ + (name, start + 1, start + len(chunk), + ','.join(str(int(v)) for v in chunk))) + return lines + + def get_color_init_routine(self, matrix_element, proc_prefix): + """Fortran source rebuilding the color matrix from its compressed + description, or an empty routine when the entries are written out.""" + + encoding = self.get_color_matrix_encoding(matrix_element) + nb_color = len(matrix_element.get('color_matrix')._sorted_keys1) \ + if matrix_element.get('color_matrix') else 0 + header = [" SUBROUTINE %sINIT_CF()" % proc_prefix] + if not encoding: + return header + [" RETURN", " END"] + + nb_rep = len(encoding['rows']) + nb_gen = len(encoding['gens']) + body = header + [ + "C Rebuild the color matrix from one line per", + "C orbit of the index permutations leaving the", + "C color basis invariant. Every other line is one", + "C of those with its columns permuted, which is", + "C what following CFPAR back to the representative", + "C line gives. Done once, on the first call.", + " IMPLICIT NONE", + " INTEGER NCOLOR, NCFREP, NCFGEN", + " PARAMETER (NCOLOR=%d)" % nb_color, + " PARAMETER (NCFREP=%d)" % nb_rep, + " PARAMETER (NCFGEN=%d)" % nb_gen, + " INTEGER %sCF(NCOLOR*(NCOLOR+1)/2)" % proc_prefix, + " INTEGER %sDENOM" % proc_prefix, + " COMMON /%scolor_matrix/ %sCF,%sDENOM" % \ + (proc_prefix, proc_prefix, proc_prefix), + " INTEGER CFROW(NCOLOR*NCFREP)", + " INTEGER CFGEN(NCOLOR*NCFGEN)", + " INTEGER CFPAR(2*NCOLOR)", + " INTEGER CFSLOT(NCOLOR)", + " INTEGER PERM(NCOLOR)", + " INTEGER I,J,NODE,G,CF_INDEX,BASE", + " LOGICAL CF_DONE", + " DATA CF_DONE/.FALSE./", + " SAVE CF_DONE", + ] + body += self.get_int_data_lines("CFROW", + sum(encoding['rows'], [])) + body += self.get_int_data_lines("CFGEN", + sum(([x + 1 for x in g] for g in encoding['gens']), + [])) + body += self.get_int_data_lines("CFPAR", + sum(([p[0], p[1]] for p in encoding['parent']), [])) + body += self.get_int_data_lines("CFSLOT", encoding['slot']) + body += [ + " IF (CF_DONE) RETURN", + " CF_DONE = .TRUE.", + " CF_INDEX = 0", + " DO I = 1, NCOLOR", + " DO J = 1, NCOLOR", + " PERM(J) = J", + " ENDDO", + " NODE = I", + " DO WHILE (CFPAR(2*NODE-1) .NE. 0)", + " G = (CFPAR(2*NODE)-1)*NCOLOR", + " DO J = 1, NCOLOR", + " PERM(J) = CFGEN(G+PERM(J))", + " ENDDO", + " NODE = CFPAR(2*NODE-1)", + " ENDDO", + " BASE = (CFSLOT(NODE)-1)*NCOLOR", + " CF_INDEX = CF_INDEX + 1", + " %sCF(CF_INDEX) = CFROW(BASE+PERM(I))" % proc_prefix, + " DO J = I+1, NCOLOR", + " CF_INDEX = CF_INDEX + 1", + " %sCF(CF_INDEX) = 2*CFROW(BASE+PERM(J))" % proc_prefix, + " ENDDO", + " ENDDO", + " END", + ] + return body + def get_den_factor_line(self, matrix_element): """Return the denominator factor line for this matrix element""" return "DATA IDEN/%2r/" % \ @@ -4324,6 +4475,9 @@ def write_matrix_element_v4(self, writer, matrix_element, fortran_model, # Extract color data lines color_data_lines = self.get_color_data_lines(matrix_element) replace_dict['color_data_lines'] = "\n".join(color_data_lines) % {'proc_prefix': replace_dict['proc_prefix']} + replace_dict['color_init_routine'] = "\n".join( + self.get_color_init_routine(matrix_element, + replace_dict['proc_prefix'])) if self.opt['export_format']=='standalone_msP': # For MadSpin need to return the AMP2 diff --git a/madgraph/iolibs/template_files/matrix_standalone_v4.inc b/madgraph/iolibs/template_files/matrix_standalone_v4.inc index ada2fd0eb1..0217236eca 100644 --- a/madgraph/iolibs/template_files/matrix_standalone_v4.inc +++ b/madgraph/iolibs/template_files/matrix_standalone_v4.inc @@ -369,8 +369,8 @@ C C LOCAL VARIABLES C - INTEGER I,J - COMPLEX*16 ZTEMP + INTEGER I,J,NJ,NB + COMPLEX*16 ZTEMP,Z1,Z2,Z3,Z4 INTEGER CF_INDEX INTEGER %(proc_prefix)sCF(NCOLOR*(NCOLOR+1)/2) @@ -383,19 +383,39 @@ C C COLOR DATA C + CALL %(proc_prefix)sINIT_CF() MATRIX = 0.D0 CF_INDEX = 0 +C Four accumulators, not one: with a single one every +C term waits for the one before it to come out of the +C adder, and that latency is what the loop spends its +C time on. No compiler does this by itself, since it +C changes the order the terms are summed in. DO I = 1, NCOLOR - ZTEMP = (0.D0,0.D0) - DO J = I, NCOLOR - CF_INDEX = CF_INDEX + 1 - ZTEMP = ZTEMP + %(proc_prefix)sCF(CF_INDEX)*JAMP(J) + Z1 = (0.D0,0.D0) + Z2 = (0.D0,0.D0) + Z3 = (0.D0,0.D0) + Z4 = (0.D0,0.D0) + NJ = NCOLOR - I + 1 + NB = (NJ/4)*4 + DO J = 0, NB-4, 4 + Z1 = Z1 + %(proc_prefix)sCF(CF_INDEX+J+1)*JAMP(I+J) + Z2 = Z2 + %(proc_prefix)sCF(CF_INDEX+J+2)*JAMP(I+J+1) + Z3 = Z3 + %(proc_prefix)sCF(CF_INDEX+J+3)*JAMP(I+J+2) + Z4 = Z4 + %(proc_prefix)sCF(CF_INDEX+J+4)*JAMP(I+J+3) ENDDO + ZTEMP = (Z1+Z2)+(Z3+Z4) + DO J = NB, NJ-1 + ZTEMP = ZTEMP + %(proc_prefix)sCF(CF_INDEX+J+1)*JAMP(I+J) + ENDDO + CF_INDEX = CF_INDEX + NJ MATRIX = MATRIX+ZTEMP*DCONJG(JAMP(I))/%(proc_prefix)sDENOM ENDDO END +%(color_init_routine)s + SUBROUTINE %(proc_prefix)sGET_INTER(JAMP_1,JAMP_2, INTER) @@ -416,6 +436,7 @@ CF2PY INTENT(IN) :: JAMP_2 C COLOR DATA C + CALL %(proc_prefix)sINIT_CF() INTER = (0.D0,0.D0) CF_INDEX = 0 diff --git a/tests/input_files/IOTestsComparison/IOExportV4IOTest/export_matrix_element_v4_standalone/matrix.f b/tests/input_files/IOTestsComparison/IOExportV4IOTest/export_matrix_element_v4_standalone/matrix.f index 330526878a..651958aa1b 100644 --- a/tests/input_files/IOTestsComparison/IOExportV4IOTest/export_matrix_element_v4_standalone/matrix.f +++ b/tests/input_files/IOTestsComparison/IOExportV4IOTest/export_matrix_element_v4_standalone/matrix.f @@ -488,8 +488,8 @@ SUBROUTINE GET_MATRIX(JAMP,MATRIX) C LOCAL VARIABLES C - INTEGER I,J - COMPLEX*16 ZTEMP + INTEGER I,J,NJ,NB + COMPLEX*16 ZTEMP,Z1,Z2,Z3,Z4 INTEGER CF_INDEX INTEGER CF(NCOLOR*(NCOLOR+1)/2) @@ -502,19 +502,41 @@ SUBROUTINE GET_MATRIX(JAMP,MATRIX) C COLOR DATA C + CALL INIT_CF() MATRIX = 0.D0 CF_INDEX = 0 +C Four accumulators, not one: with a single one every +C term waits for the one before it to come out of the +C adder, and that latency is what the loop spends its +C time on. No compiler does this by itself, since it +C changes the order the terms are summed in. DO I = 1, NCOLOR - ZTEMP = (0.D0,0.D0) - DO J = I, NCOLOR - CF_INDEX = CF_INDEX + 1 - ZTEMP = ZTEMP + CF(CF_INDEX)*JAMP(J) + Z1 = (0.D0,0.D0) + Z2 = (0.D0,0.D0) + Z3 = (0.D0,0.D0) + Z4 = (0.D0,0.D0) + NJ = NCOLOR - I + 1 + NB = (NJ/4)*4 + DO J = 0, NB-4, 4 + Z1 = Z1 + CF(CF_INDEX+J+1)*JAMP(I+J) + Z2 = Z2 + CF(CF_INDEX+J+2)*JAMP(I+J+1) + Z3 = Z3 + CF(CF_INDEX+J+3)*JAMP(I+J+2) + Z4 = Z4 + CF(CF_INDEX+J+4)*JAMP(I+J+3) + ENDDO + ZTEMP = (Z1+Z2)+(Z3+Z4) + DO J = NB, NJ-1 + ZTEMP = ZTEMP + CF(CF_INDEX+J+1)*JAMP(I+J) ENDDO + CF_INDEX = CF_INDEX + NJ MATRIX = MATRIX+ZTEMP*DCONJG(JAMP(I))/DENOM ENDDO END + SUBROUTINE INIT_CF() + RETURN + END + SUBROUTINE GET_INTER(JAMP_1,JAMP_2, INTER) @@ -535,6 +557,7 @@ SUBROUTINE GET_INTER(JAMP_1,JAMP_2, INTER) C COLOR DATA C + CALL INIT_CF() INTER = (0.D0,0.D0) CF_INDEX = 0 diff --git a/tests/unit_tests/core/test_color_amp.py b/tests/unit_tests/core/test_color_amp.py index c0dd0c7719..b507eb7987 100755 --- a/tests/unit_tests/core/test_color_amp.py +++ b/tests/unit_tests/core/test_color_amp.py @@ -945,3 +945,55 @@ def test_color_factor_simplify_merges_like_strings(self): color.ColorString([color.T(1, 2, 3)], coeff=fractions.Fraction(-1, 3))]) self.assertEqual(len(col_fact.simplify()), 0) + + def test_spanning_generators(self): + """The spanning subset has to still reach every line, and dropping the + redundant generators has to make a real difference.""" + + col_basis = color_amp.ColorBasis(self.get_gluon_amplitude(3)) + keys = sorted(col_basis.keys()) + symmetry = color_amp.ColorBasisSymmetry(keys) + kept = symmetry.spanning_generators() + self.assertTrue(0 < len(kept) < len(symmetry.generators1)) + + representatives, representative, parent, gens = symmetry.spanning_tree() + self.assertEqual(len(gens), len(kept)) + # a single orbit, and every line but the representative has a parent + self.assertEqual(representatives, [0]) + self.assertEqual(representative, [0] * len(keys)) + self.assertEqual(parent[0], None) + self.assertTrue(all(parent[i] is not None + for i in range(1, len(keys)))) + # every parent link is a generator applied to the parent line + for line in range(1, len(keys)): + origin, local = parent[line] + self.assertEqual(gens[local][origin], line) + + def test_spanning_tree_rebuilds_color_matrix(self): + """Following the parents back to the representative gives the column + permutation relating the two lines, which has to rebuild the matrix.""" + + for amplitude in [self.get_gluon_amplitude(2), + self.get_gluon_amplitude(3), + self.get_quark_amplitude(1)]: + col_basis = color_amp.ColorBasis(amplitude) + col_matrix = color_amp.ColorMatrix(col_basis, Nc=3) + keys = sorted(col_basis.keys()) + symmetry = color_amp.ColorBasisSymmetry(keys) + representatives, representative, parent, gens = \ + symmetry.spanning_tree() + denominator = max(col_matrix.get_line_denominators()) + rows = [col_matrix.get_line_numerators(i, denominator) + for i in range(len(keys))] + + for line in range(len(keys)): + # walk up to the representative, permuting the columns + perm = list(range(len(keys))) + node = line + while parent[node] is not None: + origin, local = parent[node] + perm = [gens[local][p] for p in perm] + node = origin + self.assertEqual(node, representative[line]) + self.assertEqual([rows[node][perm[j]] + for j in range(len(keys))], rows[line]) From b66e13ec00d5ed1d1e2453480427e018a75eb4a2 Mon Sep 17 00:00:00 2001 From: Olivier Mattelaer Date: Thu, 6 Aug 2026 11:55:37 +0200 Subject: [PATCH 05/42] look for the JAMP sub-expressions by orbits of the color basis symmetry A permutation of the external color indices which maps the color basis onto itself also permutes the columns of the JAMP matrix up to a sign, so the whole matrix is invariant and the sub-expressions the optimisation looks for come in orbits: every one of them is reused exactly as often as the others. optimise_jamp took them one at a time, in whatever order the argmax set happened to be in, so its result was not closed under the symmetry: only about a fifth of the definitions had their image among the definitions. Taking a whole orbit at a time instead keeps the matrix invariant at every step. Two sub-expressions of the same orbit never want the same entry of the matrix, and the contention between orbits is settled by applying one orbit at a time. That also compresses better, not worse, everywhere it was measured, since the orbit which is applied is no longer eaten into by the ones applied before it: g g > 3g 72 -> 66 operations, 4g 951 -> 795, 5g 22221 -> 15750, 6g 441943 -> 237510, t t~ g g 281 -> 264, u u~ > u u~ g g 153 -> 148, g g > u u~ d d~ g 1753 -> 1680. Only one line per orbit is then written out, plus the amplitude permutations, and INIT_JAMP walks each orbit once on the first call to work out the operands of every definition. For g g > 6g that is 157 recipes and 6 permutations instead of 441943 lines: matrix.f goes from 43.4 MB to 19.4 MB and the GET_JAMP block from 33.6 MB to 3.69 MB. The operands are read from one array holding the amplitudes first and the definitions after them, which is why AMP is declared longer. Below five thousand definitions the lines are still both smaller and faster written out, which is where the threshold comes from. Co-Authored-By: Claude Opus 5 --- madgraph/iolibs/export_v4.py | 823 +++++++++++++++++- .../template_files/matrix_standalone_v4.inc | 14 +- .../matrix.f | 12 +- tests/unit_tests/iolibs/test_export_v4.py | 186 ++++ 4 files changed, 993 insertions(+), 42 deletions(-) diff --git a/madgraph/iolibs/export_v4.py b/madgraph/iolibs/export_v4.py index 777cd3427d..225ee0dafb 100755 --- a/madgraph/iolibs/export_v4.py +++ b/madgraph/iolibs/export_v4.py @@ -227,6 +227,16 @@ class ProcessExporterFortran(VirtualExporter): } grouped_mode = False jamp_optim = False + # write the JAMP definitions as one recipe per orbit of the permutations + # leaving the color basis invariant, instead of one line per definition + jamp_orbit = False + # Below this many definitions writing them out is both smaller and faster: + # the lines still fit in the instruction cache, while the loop reading the + # operands from a table pays for the two indirections whatever the size. + # Measured on g g > n g, the two cost the same at about five thousand + # definitions (795 definitions: 0.47 us written out against 1.40 us; + # 9990 definitions: 39.4 us against 26.4 us). + jamp_orbit_min_def = 5000 # how much smaller the compressed color matrix has to be before it is used # instead of writing every entry out (see get_color_matrix_encoding) color_encoding_margin = 4 @@ -2537,14 +2547,17 @@ def get_JAMP_lines_split_order(self, col_amps, split_order_amps, return res_list, max_tmp - def get_JAMP_lines(self, col_amps, JAMP_format="JAMP(%s)", AMP_format="AMP(%s)", - split=-1): - """Return the JAMP = sum(fermionfactor * AMP(i)) lines from col_amps + def get_JAMP_lines(self, col_amps, JAMP_format="JAMP(%s)", AMP_format="AMP(%s)", + split=-1, orbit=False, proc_prefix=''): + """Return the JAMP = sum(fermionfactor * AMP(i)) lines from col_amps defined as a matrix element or directly as a color_amplitudes dictionary, - Jamp_formatLC should be define to allow to add LeadingColor computation + Jamp_formatLC should be define to allow to add LeadingColor computation (usefull for MatchBox) The split argument defines how the JAMP lines should be split in order - not to be too long.""" + not to be too long. + With orbit on, the common sub-expressions are looked for in a way which + respects the permutations leaving the color basis invariant, so that + they can be written as one recipe per orbit (see optimise_jamp).""" # Let the user call get_JAMP_lines directly from a MatrixElement or from # the color amplitudes lists. @@ -2643,11 +2656,14 @@ def get_JAMP_lines(self, col_amps, JAMP_format="JAMP(%s)", AMP_format="AMP(%s)", start_time = 0 res_list = [] - + self.myjamp_count = 0 for key in all_element: all_element[key] = complex(all_element[key]) - new_mat, defs = self.optimise_jamp(all_element) + self.jamp_orbits = None + symmetry = self.get_jamp_symmetry(col_amps, all_element) \ + if orbit and self.jamp_orbit else None + new_mat, defs = self.optimise_jamp(all_element, symmetry=symmetry) if start_time: logger.info("Color-Flow passed to %s term in %ss. Introduce %i contraction", len(new_mat), int(time.time()-start_time), len(defs)) @@ -2669,22 +2685,47 @@ def format(frac): - for i, amp1, amp2, frac, nb in defs: - if amp1 > 0: - amp1 = AMP_format % amp1 - else: - amp1 = "TMP_JAMP(%d)" % -amp1 - if amp2 > 0: - amp2 = AMP_format % amp2 - else: - amp2 = "TMP_JAMP(%d)" % -amp2 - - if frac not in [1., -1]: - res_list.append(' TMP_JAMP(%d) = %s + (%s) * %s ! used %d times' % (i,amp1, format(frac), amp2, nb)) - elif frac == 1.: - res_list.append(' TMP_JAMP(%d) = %s + %s ! used %d times' % (i,amp1, amp2, nb)) - else: - res_list.append(' TMP_JAMP(%d) = %s - %s ! used %d times' % (i,amp1, amp2, nb)) + # One recipe per orbit rather than one line per definition, when the + # symmetry allows it and there are enough definitions for the routine + # rebuilding them to be worth its own code. + recipes = None + if symmetry and len(defs) >= self.jamp_orbit_min_def: + recipes = self.jamp_orbit_recipes(defs, + col_amps.get_number_of_amplitudes()) + self.jamp_recipes = recipes + + if recipes: + tmp_name = lambda k: "AMP(NGRAPHS+%d)" % k + defs = recipes['defs'] + res_list.append("C The definitions below come in orbits of the") + res_list.append("C permutations leaving the color basis") + res_list.append("C invariant: all of an orbit are the same") + res_list.append("C recipe with the amplitudes permuted, so") + res_list.append("C only their operands differ. INIT_JAMP works") + res_list.append("C those out once, from one recipe per orbit.") + res_list.append(" CALL %sINIT_JAMP()" % proc_prefix) + res_list.append(" DO ITMP = 1, NB_TMP_JAMP") + res_list.append(" AMP(NGRAPHS+ITMP) = AMP(TMP_JAMP_A(ITMP))" + " + TMP_JAMP_F(ITMP)*AMP(TMP_JAMP_B(ITMP))") + res_list.append(" ENDDO") + else: + tmp_name = lambda k: "TMP_JAMP(%d)" % k + for i, amp1, amp2, frac, nb in defs: + if amp1 > 0: + amp1 = AMP_format % amp1 + else: + amp1 = tmp_name(-amp1) + if amp2 > 0: + amp2 = AMP_format % amp2 + else: + amp2 = tmp_name(-amp2) + + if frac not in [1., -1]: + res_list.append(' TMP_JAMP(%d) = %s + (%s) * %s ! used %d times' % (i,amp1, format(frac), amp2, nb)) + elif frac == 1.: + res_list.append(' TMP_JAMP(%d) = %s + %s ! used %d times' % (i,amp1, amp2, nb)) + else: + res_list.append(' TMP_JAMP(%d) = %s - %s ! used %d times' % (i,amp1, amp2, nb)) jamp_res = collections.defaultdict(list) max_jamp=0 @@ -2692,7 +2733,12 @@ def format(frac): if var > 0: name = AMP_format % var else: - name = "TMP_JAMP(%d)" % -var + if recipes: + # the definitions were renumbered, and one of them can be + # the opposite of the one the optimisation had + where, scale = recipes['factor_of'][-var] + factor, var = factor * scale, -where + name = tmp_name(-var) if factor not in [1.]: jamp_res[jamp].append("(%s)*%s" % (format(factor), name)) elif factor ==1: @@ -2748,13 +2794,22 @@ def common_jamp_lines(columns, nb_line, j1, j2): pos2 += 1 return res - def optimise_jamp(self, all_element, nb_line=0, nb_col=0, added=0): + def optimise_jamp(self, all_element, nb_line=0, nb_col=0, added=0, + symmetry=None): """ optimise problem of type Y = A X A is a matrix (all_element) X is the fortran name of the input. The code iteratively add sub-expression jtemp[sub_add] and recall itself (this is add to the X size) + + With a symmetry (see get_jamp_symmetry) the sub-expressions are + introduced by whole orbits of that symmetry instead of one at a + time, so that the result can be written as one recipe per orbit. + The orbits are then left in self.jamp_orbits. """ + if symmetry: + return self.optimise_jamp_equivariant(all_element, symmetry) + self.myjamp_count +=1 if not nb_line: @@ -2853,11 +2908,687 @@ def optimise_jamp(self, all_element, nb_line=0, nb_col=0, added=0): new_element, new_def = self.optimise_jamp(all_element, nb_line=nb_line, nb_col=nb_col, added=added) for one_def in to_add: new_def.insert(0, one_def) - return new_element, new_def - - - - + return new_element, new_def + + + #=========================================================================== + # Orbit equivariant version of the JAMP optimisation + #=========================================================================== + # A permutation of the external color indices which maps the color basis + # onto itself (see color_amp.ColorBasisSymmetry) also permutes the columns + # of the JAMP matrix, up to a sign. The whole matrix is then invariant, so + # the sub-expressions the optimisation looks for come in orbits: every one + # of them is worth exactly as much as the others. Introducing a whole orbit + # at a time, rather than one sub-expression at a time as the plain scan + # does, leaves the matrix invariant at every step, and the definitions can + # be written as one recipe per orbit. + + @staticmethod + def jamp_column_form(column): + """Canonical form of one column of the JAMP matrix up to a global sign, + together with the sign which was taken out.""" + + entries = sorted(column.items()) + first = entries[0][1] + sign = -1 if (first.real, first.imag) < (0., 0.) else 1 + return tuple((i, sign * value) for i, value in entries), sign + + @classmethod + def jamp_amp_permutation(cls, columns, induced): + """Permutation of the amplitudes induced by the permutation induced of + the color basis: return {amp: (amp, sign)} such that + + M[induced[i], sigma(j)] = sign(j) * M[i, j] + + or None if the columns are not mapped onto each other. + + Several amplitudes often have the very same column, so the columns are + gathered by their canonical form and one target is taken out of each + group at a time: looking the image up would not give a bijection.""" + + groups = collections.defaultdict(collections.deque) + for j in sorted(columns): + form, sign = cls.jamp_column_form(columns[j]) + groups[form].append((j, sign)) + + action = {} + for j in sorted(columns): + image = dict((induced[i - 1] + 1, value) + for i, value in columns[j].items()) + form, sign = cls.jamp_column_form(image) + group = groups.get(form) + if not group: + return None + target, target_sign = group.popleft() + factor = sign * target_sign + other = columns[target] + if len(other) != len(image) or \ + any(other.get(i) != factor * value + for i, value in image.items()): + return None + action[j] = (target, factor) + return action + + def get_jamp_symmetry(self, matrix_element, all_element): + """Permutations leaving the JAMP matrix invariant: for each of them the + permutation of the color basis lines, and the permutation of the + amplitude columns with the sign that goes with it. None when there is + none, or when the matrix element does not carry a color basis.""" + + if not isinstance(matrix_element, helas_objects.HelasMatrixElement): + return None + color_basis = matrix_element.get('color_basis') + if not color_basis or len(color_basis) < 2: + return None + symmetry = color_amp.ColorBasisSymmetry(sorted(color_basis.keys())) + if not symmetry.generators1: + return None + + columns = collections.defaultdict(dict) + for (i, j), value in all_element.items(): + if value: + columns[j][i] = value + if not columns: + return None + + nb_line = len(symmetry.keys1) + rowperms, actions = [], [] + for induced in symmetry.generators1: + action = self.jamp_amp_permutation(columns, induced) + if action is None: + continue + rowperms.append([0] + [induced[i] + 1 for i in range(nb_line)]) + actions.append(action) + if not actions: + return None + + # one line per orbit is enough to see every sub-expression: any other + # line is the image of one of them, and so are the sub-expressions it + # holds. This is what keeps the scan below from being quadratic in the + # number of terms of the whole matrix. + parent = list(range(nb_line + 1)) + + def find(x): + while parent[x] != x: + parent[x] = parent[parent[x]] + x = parent[x] + return x + + for rowperm in rowperms: + for i in range(1, nb_line + 1): + ri, rj = find(i), find(rowperm[i]) + if ri != rj: + parent[ri] = rj + line_reps = [i for i in range(1, nb_line + 1) if find(i) == i] + + return {'rowperms': rowperms, 'actions': actions, + 'nb_line': nb_line, 'line_reps': line_reps} + + @staticmethod + def jamp_operation_image(action, operation): + """Image of the sub-expression operation=(j1,j2,R) under one + permutation, and the factor relating the column the image defines to + the image of the column operation defines.""" + + j1, j2, ratio = operation + first, sign1 = action[j1] + second, sign2 = action[j2] + if first < second: + return (first, second, ratio * sign2 / sign1), sign1 + return (second, first, sign1 / (sign2 * ratio)), sign2 * ratio + + def optimise_jamp_equivariant(self, all_element, symmetry): + """Same optimisation as optimise_jamp, but introducing whole orbits of + sub-expressions at a time so that the result is closed under the + symmetry. Fills self.jamp_orbits with, for every definition, the orbit + it belongs to and the definition and permutation it comes from.""" + + actions = [dict(action) for action in symmetry['actions']] + line_reps = symmetry['line_reps'] + added = 0 + defs = [] + # (orbit, parent definition, permutation) for every definition + tree = [] + nb_orbit = 0 + + while True: + columns = collections.defaultdict(list) + lines = collections.defaultdict(list) + for (i, j), value in all_element.items(): + if value: + columns[j].append(i) + lines[i].append(j) + for line in lines.values(): + line.sort() + + # every sub-expression is the image of one living on a + # representative line, so only those have to be looked at + candidates = set() + for i in line_reps: + line = lines.get(i, []) + for pos, j1 in enumerate(line): + value = all_element[(i, j1)] + for j2 in line[pos + 1:]: + candidates.add((j1, j2, all_element[(i, j2)] / value)) + + max_count = 0 + best = [] + for operation in candidates: + count = len(self.jamp_operation_lines(all_element, columns, + operation)) + if count > max_count: + max_count, best = count, [operation] + elif count == max_count: + best.append(operation) + if max_count <= 1: + break + + orbits = self.jamp_operation_orbits(actions, best) + first_of_level = added + 1 + for orbit, parent in orbits: + rows = dict((operation, + self.jamp_operation_lines(all_element, columns, + operation)) + for operation in orbit) + if not self.jamp_orbit_usable(rows): + continue + index = {} + for operation in orbit: + added += 1 + index[operation] = added + origin, permutation = parent[operation] + tree.append((nb_orbit, index[origin] if origin else 0, + permutation)) + defs.append((added, operation[0], operation[1], + operation[2], len(rows[operation]))) + nb_orbit += 1 + for operation, new in index.items(): + j1, j2 = operation[0], operation[1] + for i in rows[operation]: + all_element[(i, -new)] = all_element[(i, j1)] + del all_element[(i, j1)] + del all_element[(i, j2)] + for action in actions: + for operation, new in index.items(): + image, factor = self.jamp_operation_image(action, + operation) + action[-new] = (-index[image], factor) + if added < first_of_level: + # nothing could be introduced as a whole orbit + break + logger.log(5, "Define %d new shortcut reused %d times", + added - first_of_level + 1, max_count) + + self.jamp_orbits = {'tree': tree, 'nb_orbit': nb_orbit, + 'actions': actions, 'symmetry': symmetry} + return all_element, defs + + @staticmethod + def jamp_operation_lines(all_element, columns, operation): + """Lines where both columns of the sub-expression are still there with + its ratio. The values are read from the matrix as it is now, so lines + already taken by an orbit introduced before are simply gone.""" + + j1, j2, ratio = operation + res = [] + for i in columns.get(j1, ()): + value = all_element.get((i, j1), 0) + if not value: + continue + other = all_element.get((i, j2), 0) + if other and other / value == ratio: + res.append(i) + return res + + def jamp_operation_orbits(self, actions, operations): + """Orbits of the sub-expressions, walked breadth first, with the + (sub-expression, permutation) each of them is reached from.""" + + seen = set() + orbits = [] + for start in sorted(operations, key=lambda op: (op[0], op[1], + op[2].real, + op[2].imag)): + if start in seen: + continue + orbit, parent = [start], {start: (None, 0)} + seen.add(start) + queue = collections.deque([start]) + while queue: + current = queue.popleft() + for position, action in enumerate(actions): + image = self.jamp_operation_image(action, current)[0] + if image in seen: + continue + seen.add(image) + parent[image] = (current, position + 1) + orbit.append(image) + queue.append(image) + orbits.append((orbit, parent)) + return orbits + + def jamp_orbit_recipes(self, defs, nb_amp): + """Describe the definitions by one recipe per orbit: the amplitude + permutations, the first definition of every orbit, and the definitions + renumbered so that walking each orbit breadth first from its recipe, + with those permutations in that order, hands them out in that very + order. The generated code walks them the same way, so it only needs + the recipes. + + Returns None when the definitions cannot be described this way, and + the caller then writes them out one by one as before.""" + + orbits = self.jamp_orbits + if not orbits or not defs: + return None + actions = orbits['actions'] + nb_orbit = orbits['nb_orbit'] + + # only a plain sign in front of a definition can be carried by the + # index of that definition alone, which is what keeps the generated + # routine to integer arithmetic + for one_def in defs: + if one_def[3] not in (1, -1): + return None + for action in actions: + for column, (_image, factor) in action.items(): + if column < 0 and factor not in (1, -1): + return None + + # first definition of every orbit + first = [0] * nb_orbit + for (orbit, parent, _permutation), one_def in zip(orbits['tree'], defs): + if not parent: + first[orbit] = one_def[0] + + # the permutations which are really needed to reach every definition + # of every orbit: each of them costs one table of amplitude indices + chosen = [] + rest = list(range(len(actions))) + while self.jamp_orbit_reach(actions, chosen, first) < len(defs): + best, best_gain = None, -1 + for position in rest: + gain = self.jamp_orbit_reach(actions, chosen + [position], + first) + if gain > best_gain: + best, best_gain = position, gain + if best is None: + return None + chosen.append(best) + rest.remove(best) + + replay = self.jamp_orbit_replay(defs, first, chosen) + while replay is None and rest: + # the permutations kept do not reach every definition after all + chosen.append(rest.pop(0)) + replay = self.jamp_orbit_replay(defs, first, chosen) + if replay is None: + return None + new_defs, recipes, factor_of = replay + + permutations = [] + for permutation in chosen: + action = orbits['symmetry']['actions'][permutation] + row = [0] * nb_amp + for amp, (image, sign) in action.items(): + row[amp - 1] = image if sign > 0 else -image + if any(value == 0 for value in row): + return None + permutations.append(row) + + return {'permutations': permutations, 'recipes': recipes, + 'defs': new_defs, 'nb_amp': nb_amp, 'factor_of': factor_of} + + @staticmethod + def jamp_hash_size(nb_def): + """A prime comfortably larger than twice the number of definitions: + the routine which rebuilds them looks the operand pairs up in a table + of that size with linear probing.""" + + candidate = 2 * nb_def + 101 + while True: + candidate += 1 + for divisor in range(2, int(candidate ** 0.5) + 1): + if candidate % divisor == 0: + break + else: + return candidate + + def get_jamp_decl_lines(self, recipes, proc_prefix): + """The declarations GET_JAMP needs to run the definitions.""" + + if not recipes: + return [] + return [ + " INTEGER ITMP", + " INTEGER NB_TMP_JAMP", + " PARAMETER (NB_TMP_JAMP=%d)" % len(recipes['defs']), + " INTEGER TMP_JAMP_A(NB_TMP_JAMP), TMP_JAMP_B(NB_TMP_JAMP)", + " DOUBLE PRECISION TMP_JAMP_F(NB_TMP_JAMP)", + " COMMON /%sjamp_recipe/ TMP_JAMP_A,TMP_JAMP_B,TMP_JAMP_F" % \ + proc_prefix, + ] + + def get_jamp_init_routine(self, recipes, proc_prefix): + """Fortran source rebuilding the operands of every color flow + definition from one recipe per orbit, or nothing when the definitions + are written out.""" + + if not recipes: + return [] + nb_def = len(recipes['defs']) + nb_amp = recipes['nb_amp'] + nb_orbit = len(recipes['recipes']) + nb_perm = len(recipes['permutations']) + nb_hash = self.jamp_hash_size(nb_def) + + common = [ + " INTEGER NGRAPHS, NB_TMP_JAMP, NB_HASH", + " PARAMETER (NGRAPHS=%d)" % nb_amp, + " PARAMETER (NB_TMP_JAMP=%d)" % nb_def, + " PARAMETER (NB_HASH=%d)" % nb_hash, + " INTEGER TMP_JAMP_A(NB_TMP_JAMP), TMP_JAMP_B(NB_TMP_JAMP)", + " DOUBLE PRECISION TMP_JAMP_F(NB_TMP_JAMP)", + " COMMON /%sjamp_recipe/ TMP_JAMP_A,TMP_JAMP_B,TMP_JAMP_F" % \ + proc_prefix, + " INTEGER NUSED", + " INTEGER HVAL(NB_HASH)", + " INTEGER*8 HKEY(NB_HASH)", + " COMMON /%sjamp_build/ NUSED,HVAL,HKEY" % proc_prefix, + ] + + add = [ + " SUBROUTINE %sJAMP_ADD(A,B,F,M,SWAP)" % proc_prefix, + "C The definition A + F*B, added if it is not there yet.", + "C Its two operands the other way round give the very same", + "C column times F, so that one is looked for as well and SWAP", + "C says which of the two was found.", + " IMPLICIT NONE", + " INTEGER A,B,F,M,SWAP", + ] + common + [ + " INTEGER H, FREE, SHIFT", + " INTEGER*8 KEY, OTHER, BASE", + " SHIFT = NB_TMP_JAMP + 1", + " BASE = NGRAPHS + NB_TMP_JAMP + 2", + " KEY = ((A+SHIFT)*BASE+(B+SHIFT))*2+(1-F)/2+1", + " OTHER = ((B+SHIFT)*BASE+(A+SHIFT))*2+(1-F)/2+1", + " SWAP = 1", + " H = INT(MOD(KEY,INT(NB_HASH,8)))+1", + " DO WHILE (HKEY(H) .NE. 0)", + " IF (HKEY(H) .EQ. KEY) THEN", + " M = HVAL(H)", + " RETURN", + " ENDIF", + " H = H+1", + " IF (H .GT. NB_HASH) H = 1", + " ENDDO", + " FREE = H", + " H = INT(MOD(OTHER,INT(NB_HASH,8)))+1", + " DO WHILE (HKEY(H) .NE. 0)", + " IF (HKEY(H) .EQ. OTHER) THEN", + " M = HVAL(H)", + " SWAP = F", + " RETURN", + " ENDIF", + " H = H+1", + " IF (H .GT. NB_HASH) H = 1", + " ENDDO", + " NUSED = NUSED+1", + " M = NUSED", + " TMP_JAMP_A(M) = A", + " TMP_JAMP_B(M) = B", + " TMP_JAMP_F(M) = F", + " HKEY(FREE) = KEY", + " HVAL(FREE) = M", + " END", + "", + ] + + body = [ + " SUBROUTINE %sINIT_JAMP()" % proc_prefix, + "C Work out the operands of every color flow definition,", + "C starting from one recipe per orbit of the permutations", + "C leaving the color basis invariant and walking each orbit", + "C with those permutations. Done once, on the first call.", + " IMPLICIT NONE", + " INTEGER NB_ORBIT, NB_PERM", + " PARAMETER (NB_ORBIT=%d)" % nb_orbit, + " PARAMETER (NB_PERM=%d)" % nb_perm, + ] + common + [ + " INTEGER JPERM(NGRAPHS*NB_PERM)", + " INTEGER JREC(3*NB_ORBIT)", + " INTEGER JIMG(NB_TMP_JAMP*NB_PERM)", + " INTEGER I,J,P,A,B,T,SA,SB,M,SWAP,BEGIN", + " LOGICAL JAMP_DONE", + " DATA JAMP_DONE/.FALSE./", + " SAVE JAMP_DONE, JIMG", + ] + body += self.get_int_data_lines("JPERM", + sum(recipes['permutations'], [])) + body += self.get_int_data_lines("JREC", + sum((list(one) + for one in recipes['recipes']), + [])) + body += [ + " IF (JAMP_DONE) RETURN", + " JAMP_DONE = .TRUE.", + " DO I = 1, NB_HASH", + " HKEY(I) = 0", + " ENDDO", + " NUSED = 0", + " DO I = 1, NB_ORBIT", + " BEGIN = NUSED", + " CALL %sJAMP_ADD(JREC(3*I-2),JREC(3*I-1),JREC(3*I),M,SWAP)" + % proc_prefix, + " J = BEGIN", + " DO WHILE (J .LT. NUSED)", + " J = J+1", + " DO P = 1, NB_PERM", + " A = TMP_JAMP_A(J)", + " IF (A .GT. 0) THEN", + " T = JPERM((P-1)*NGRAPHS+A)", + " A = ABS(T)", + " ELSE", + " T = JIMG((-A-1)*NB_PERM+P)", + " A = -ABS(T)", + " ENDIF", + " SA = ISIGN(1,T)", + " B = TMP_JAMP_B(J)", + " IF (B .GT. 0) THEN", + " T = JPERM((P-1)*NGRAPHS+B)", + " B = ABS(T)", + " ELSE", + " T = JIMG((-B-1)*NB_PERM+P)", + " B = -ABS(T)", + " ENDIF", + " SB = ISIGN(1,T)", + " T = SA*SB*NINT(TMP_JAMP_F(J))", + " CALL %sJAMP_ADD(A,B,T,M,SWAP)" % proc_prefix, + " JIMG((J-1)*NB_PERM+P) = SA*SWAP*M", + " ENDDO", + " ENDDO", + " ENDDO", + " IF (NUSED .NE. NB_TMP_JAMP) THEN", + " WRITE(*,*) 'ERROR: color flow recipes gave',NUSED,", + " $ ' definitions instead of',NB_TMP_JAMP", + " STOP 1", + " ENDIF", + "C the operands are read from one array holding the", + "C amplitudes first and the definitions after them", + " DO I = 1, NB_TMP_JAMP", + " IF (TMP_JAMP_A(I) .LT. 0) TMP_JAMP_A(I) = NGRAPHS" + "-TMP_JAMP_A(I)", + " IF (TMP_JAMP_B(I) .LT. 0) TMP_JAMP_B(I) = NGRAPHS" + "-TMP_JAMP_B(I)", + " ENDDO", + " END", + ] + return add + body + + def jamp_orbit_allowed(self, matrix_element): + """The orbit recipes need the routine which rebuilds the definitions + at run time, which only the plain standalone template carries.""" + + if not self.jamp_orbit or type(self) is not ProcessExporterFortranSA: + return False + if self.matrix_template != 'matrix_standalone_v4.inc': + return False + if self.opt.get('export_format') in ('standalone_msP', + 'standalone_msF', 'matchbox', + 'madloop_matchbox'): + return False + return not matrix_element.get('processes')[0].get('split_orders') + + def jamp_orbit_replay(self, defs, first, chosen): + """Walk every orbit from its first definition with the given + permutations, exactly as the generated routine does, and hand out the + definition numbers in that order. Returns the definitions in the new + numbering, the recipe of every orbit, and for each old definition the + new one it became with the factor between the two. None if the walk + does not reach every definition.""" + + amp_action = [self.jamp_orbits['symmetry']['actions'][position] + for position in chosen] + nb_perm = len(chosen) + by_index = dict((one_def[0], one_def) for one_def in defs) + # old definition -> (new definition, factor between the two columns) + factor_of = {} + left_of, right_of, ratio_of, image_of = [], [], [], [] + known = {} + recipes = [] + + def store(left, right, ratio): + """The definition with those operands, added if it is new. The two + operands can also be the other way round, and the column is then + the same one up to the ratio, hence the second look up.""" + + found = known.get((left, right, ratio)) + if found is not None: + return found, 1 + found = known.get((right, left, ratio)) + if found is not None: + return found, ratio + left_of.append(left) + right_of.append(right) + ratio_of.append(ratio) + image_of.extend([0] * nb_perm) + known[(left, right, ratio)] = len(left_of) + return len(left_of), 1 + + def act(place, column): + """image of a column and the sign that goes with it""" + + if column > 0: + return amp_action[place][column] + signed = image_of[(-column - 1) * nb_perm + place] + return -abs(signed), 1 if signed > 0 else -1 + + # the same walk is followed on the definitions of the optimisation, so + # that each of them is matched with the one the generated code builds + actions = self.jamp_orbits['actions'] + origin_of = [] + + for start in first: + _k, left, right, ratio = by_index[start][:4] + scale = 1 + if left < 0: + if -left not in factor_of: + return None + new, factor = factor_of[-left] + left, scale = -new, factor + if right < 0: + if -right not in factor_of: + return None + new, factor = factor_of[-right] + right, ratio = -new, ratio * factor + ratio = ratio / scale + if ratio not in (1, -1): + return None + begin = len(left_of) + new, factor = store(left, right, ratio) + if new != begin + 1: + # the first definition of an orbit has to be a new one + return None + origin_of.append(start) + factor_of[start] = (new, scale * factor) + recipes.append((left, right, int(complex(ratio).real))) + + current = begin + while current < len(left_of): + current += 1 + previous = origin_of[current - 1] + for place in range(nb_perm): + image_left, sign_left = act(place, left_of[current - 1]) + image_right, sign_right = act(place, right_of[current - 1]) + image_ratio = ratio_of[current - 1] * sign_left * sign_right + if image_ratio not in (1, -1): + return None + where, swap = store(image_left, image_right, image_ratio) + sign = sign_left * swap + image_of[(current - 1) * nb_perm + place] = \ + where if complex(sign).real > 0 else -where + if where > len(origin_of): + origin_of.append(None) + # follow the same step on the definitions of the + # optimisation to know which one this is + image, factor = actions[chosen[place]][-previous] + image = -image + if origin_of[where - 1] is None: + origin_of[where - 1] = image + if image not in factor_of: + factor_of[image] = (where, + factor_of[previous][1] * sign + / factor) + + if len(factor_of) != len(defs): + return None + + new_defs = [(i + 1, left_of[i], right_of[i], ratio_of[i], 0) + for i in range(len(left_of))] + return new_defs, recipes, factor_of + + @staticmethod + def jamp_orbit_reach(actions, chosen, first): + """How many definitions the given permutations reach from the first + definition of every orbit.""" + + seen = set(first) + queue = collections.deque(first) + while queue: + current = queue.popleft() + for permutation in chosen: + image = -actions[permutation][-current][0] + if image not in seen: + seen.add(image) + queue.append(image) + return len(seen) + + @staticmethod + def jamp_orbit_usable(rows): + """Restrict an orbit to the entries only one of its sub-expressions + wants, and say whether what is left can be introduced as a whole. Which + of two sub-expressions of the same orbit gets a shared entry cannot be + decided in a way that commutes with the symmetry, so those entries are + left in the matrix and get another chance in a later round.""" + + sizes = set(len(use) for use in rows.values()) + if len(sizes) != 1 or sizes == set([0]): + return False + entry = collections.Counter() + for operation, use in rows.items(): + for i in use: + entry[(i, operation[0])] += 1 + entry[(i, operation[1])] += 1 + if max(entry.values()) == 1: + return True + for operation in list(rows): + rows[operation] = [i for i in rows[operation] + if entry[(i, operation[0])] == 1 + and entry[(i, operation[1])] == 1] + sizes = set(len(use) for use in rows.values()) + return len(sizes) == 1 and sizes != set([0]) + + def get_pdf_lines(self, matrix_element, ninitial, subproc_group = False, vector=False): """Generate the PDF lines for the auto_dsig.f file""" @@ -3571,6 +4302,7 @@ class ProcessExporterFortranSA(ProcessExporterFortran): f2py_wrapper_all ="f2py_wrapper_all.inc" f2py_matrix_splitter = "f2py_splitter.py" jamp_optim = True + jamp_orbit = True default_vector_size = 0 # When True, emit per-call IAND(WF_FLAVOR_MASK/AMP_FLAVOR_MASK, # CURRENT_FLAV_BIT) guards in MATRIX so that wavefunctions and amplitudes @@ -4488,11 +5220,14 @@ def write_matrix_element_v4(self, writer, matrix_element, fortran_model, # JAMP definition, depends on the number of independent split orders split_orders=matrix_element.get('processes')[0].get('split_orders') + self.jamp_recipes = None if len(split_orders)==0: replace_dict['nSplitOrders']='' # Extract JAMP lines - jamp_lines, nb_tmp_jamp = self.get_JAMP_lines(matrix_element) + jamp_lines, nb_tmp_jamp = self.get_JAMP_lines(matrix_element, + orbit=self.jamp_orbit_allowed(matrix_element), + proc_prefix=replace_dict['proc_prefix']) # Consider the output of a dummy order 'ALL_ORDERS' for which we # set all amplitude order to weight 1 and only one squared order # contribution which is of course ALL_ORDERS=2. @@ -4535,7 +5270,23 @@ def write_matrix_element_v4(self, writer, matrix_element, fortran_model, PARAMETER (NSQSO_BORN=%d)"""%replace_dict['nSqAmpSplitOrders']) files.cp('nsqso_born.inc', '..') - replace_dict['jamp_lines'] = '\n'.join(jamp_lines) + replace_dict['jamp_lines'] = '\n'.join(jamp_lines) + + # The definitions written as one recipe per orbit are held in one + # array together with the amplitudes, so that the loop running them + # reads its two operands from the same place. + recipes = getattr(self, 'jamp_recipes', None) + replace_dict['jamp_decl'] = '\n'.join( + self.get_jamp_decl_lines(recipes, replace_dict['proc_prefix'])) + replace_dict['jamp_init_routine'] = '\n'.join( + self.get_jamp_init_routine(recipes, replace_dict['proc_prefix'])) + if recipes: + replace_dict['namp_dim'] = 'NGRAPHS+%d' % replace_dict['nb_temp_jamp'] + replace_dict['jamp_tmp_decl'] = '' + else: + replace_dict['namp_dim'] = 'NGRAPHS' + replace_dict['jamp_tmp_decl'] = \ + " COMPLEX*16 TMP_JAMP(%i)" % replace_dict['nb_temp_jamp'] matrix_template = self.matrix_template if self.opt['export_format']=='standalone_msP' : @@ -4812,9 +5563,11 @@ def finalize(self, matrix_elements, history, mg5options, flaglist, second_export def get_JAMP_lines(self, col_amps, JAMP_format="JAMP(%s)", AMP_format="AMP(%s)", split=-1, - JAMP_formatLC=None): - - """Adding leading color part of the colorflow""" + JAMP_formatLC=None, orbit=False): + + """Adding leading color part of the colorflow. The leading color part + needs the definitions written out, so the orbit recipes are not used + here.""" if not JAMP_formatLC: JAMP_formatLC= "LN%s" % JAMP_format diff --git a/madgraph/iolibs/template_files/matrix_standalone_v4.inc b/madgraph/iolibs/template_files/matrix_standalone_v4.inc index 0217236eca..d8239ac9e3 100644 --- a/madgraph/iolibs/template_files/matrix_standalone_v4.inc +++ b/madgraph/iolibs/template_files/matrix_standalone_v4.inc @@ -228,7 +228,8 @@ C INTEGER %(proc_prefix)sCF(%(ncolortriang)d) INTEGER %(proc_prefix)sDENOM common /%(proc_prefix)scolor_matrix/ %(proc_prefix)sCF,%(proc_prefix)sDENOM - COMPLEX*16 AMP(NGRAPHS), JAMP(NCOLOR), TMP_JAMP(%(nb_temp_jamp)i) + COMPLEX*16 AMP(%(namp_dim)s), JAMP(NCOLOR) +%(jamp_tmp_decl)s type(aloha) W(NWAVEFUNCS) COMPLEX*16 DUM0,DUM1 DATA DUM0, DUM1/(0d0, 0d0), (1d0, 0d0)/ @@ -344,7 +345,9 @@ C PARAMETER ( NCOLOR=%(ncolor)d) COMPLEX*16 IMAG1 PARAMETER (IMAG1=(0D0,1D0)) - COMPLEX*16 AMP(NGRAPHS), JAMP(NCOLOR), TMP_JAMP(%(nb_temp_jamp)i) + COMPLEX*16 AMP(%(namp_dim)s), JAMP(NCOLOR) +%(jamp_tmp_decl)s +%(jamp_decl)s %(jamp_lines)s END @@ -376,7 +379,8 @@ C INTEGER %(proc_prefix)sCF(NCOLOR*(NCOLOR+1)/2) INTEGER %(proc_prefix)sDENOM common /%(proc_prefix)scolor_matrix/ %(proc_prefix)sCF,%(proc_prefix)sDENOM - COMPLEX*16 JAMP(NCOLOR), TMP_JAMP(%(nb_temp_jamp)i) + COMPLEX*16 JAMP(NCOLOR) +%(jamp_tmp_decl)s COMPLEX*16 DUM0,DUM1 DATA DUM0, DUM1/(0d0, 0d0), (1d0, 0d0)/ C @@ -416,6 +420,8 @@ C changes the order the terms are summed in. %(color_init_routine)s +%(jamp_init_routine)s + SUBROUTINE %(proc_prefix)sGET_INTER(JAMP_1,JAMP_2, INTER) @@ -568,7 +574,7 @@ c PARAMETER (NCOLOR=%(ncolor)d) INTEGER IC(NEXTERNAL) - DOUBLE COMPLEX AMP(NGRAPHS) + DOUBLE COMPLEX AMP(%(namp_dim)s) DOUBLE COMPLEX, ALLOCATABLE, SAVE :: JAMP(:,:) INTEGER, SAVE :: S_NCOMB = 0 diff --git a/tests/input_files/IOTestsComparison/IOExportV4IOTest/export_matrix_element_v4_standalone/matrix.f b/tests/input_files/IOTestsComparison/IOExportV4IOTest/export_matrix_element_v4_standalone/matrix.f index 651958aa1b..b8db276219 100644 --- a/tests/input_files/IOTestsComparison/IOExportV4IOTest/export_matrix_element_v4_standalone/matrix.f +++ b/tests/input_files/IOTestsComparison/IOExportV4IOTest/export_matrix_element_v4_standalone/matrix.f @@ -275,7 +275,8 @@ REAL*8 FUNCTION MATRIX(P,NHEL,IC,FLAV_IDX) INTEGER CF(1) INTEGER DENOM COMMON /COLOR_MATRIX/ CF,DENOM - COMPLEX*16 AMP(NGRAPHS), JAMP(NCOLOR), TMP_JAMP(0) + COMPLEX*16 AMP(NGRAPHS), JAMP(NCOLOR) + COMPLEX*16 TMP_JAMP(0) TYPE(ALOHA) W(NWAVEFUNCS) COMPLEX*16 DUM0,DUM1 DATA DUM0, DUM1/(0D0, 0D0), (1D0, 0D0)/ @@ -460,7 +461,9 @@ SUBROUTINE GET_JAMP(AMP,JAMP) PARAMETER ( NCOLOR=1) COMPLEX*16 IMAG1 PARAMETER (IMAG1=(0D0,1D0)) - COMPLEX*16 AMP(NGRAPHS), JAMP(NCOLOR), TMP_JAMP(0) + COMPLEX*16 AMP(NGRAPHS), JAMP(NCOLOR) + COMPLEX*16 TMP_JAMP(0) + JAMP(1) = (-1.000000000000000D+00)*AMP(1)+(-1.000000000000000D $ +00)*AMP(2)+(-1.000000000000000D+00)*AMP(3)+( @@ -495,7 +498,8 @@ SUBROUTINE GET_MATRIX(JAMP,MATRIX) INTEGER CF(NCOLOR*(NCOLOR+1)/2) INTEGER DENOM COMMON /COLOR_MATRIX/ CF,DENOM - COMPLEX*16 JAMP(NCOLOR), TMP_JAMP(0) + COMPLEX*16 JAMP(NCOLOR) + COMPLEX*16 TMP_JAMP(0) COMPLEX*16 DUM0,DUM1 DATA DUM0, DUM1/(0D0, 0D0), (1D0, 0D0)/ C @@ -539,6 +543,8 @@ SUBROUTINE INIT_CF() + + SUBROUTINE GET_INTER(JAMP_1,JAMP_2, INTER) CF2PY INTENT(OUT) :: INTER diff --git a/tests/unit_tests/iolibs/test_export_v4.py b/tests/unit_tests/iolibs/test_export_v4.py index aae1514b2d..c7ec3a9312 100644 --- a/tests/unit_tests/iolibs/test_export_v4.py +++ b/tests/unit_tests/iolibs/test_export_v4.py @@ -10512,3 +10512,189 @@ def test_optimise_jamp_no_saving(self): result, defs = exporter.optimise_jamp(dict(all_element)) self.assertEqual(defs, []) self.assertEqual(result, all_element) + + #=========================================================================== + # Orbit equivariant optimisation + #=========================================================================== + @staticmethod + def symmetric_matrix(seed, nb_line, nb_col, density): + """A matrix which really is invariant under a permutation of its lines + and of its columns: the entries are drawn on one half of the lines and + the permutation is used to fill the other half. Returns the matrix and + the symmetry in the form optimise_jamp expects. + + The line permutation exchanges the two halves, and the column one + reverses the columns, with a sign on one pair of columns in three. The + sign only depends on the pair, so that applying the permutation twice + really gives the identity.""" + + rng = random.Random(seed) + values = [1, -1, 2, -2] + half = nb_line // 2 + + rowperm = [0] * (nb_line + 1) + for i in range(1, half + 1): + rowperm[i] = i + half + rowperm[i + half] = i + action = {} + for j in range(1, nb_col + 1): + image = nb_col + 1 - j + action[j] = (image, -1 if min(j, image) % 3 == 0 else 1) + + all_element = {} + for i in range(1, half + 1): + for j in range(1, nb_col + 1): + if rng.random() < density: + all_element[(i, j)] = complex(rng.choice(values)) + # M[rowperm[i], action[j]] = sign * M[i, j] + for (i, j), value in list(all_element.items()): + image, sign = action[j] + all_element[(rowperm[i], image)] = sign * value + + # the matrix really is invariant, otherwise the test would not be + # testing what it says it is + for (i, j), value in all_element.items(): + image, sign = action[j] + assert all_element.get((rowperm[i], image), 0) == sign * value + + symmetry = {'rowperms': [rowperm], 'actions': [action], + 'nb_line': nb_line, 'line_reps': list(range(1, half + 1))} + return all_element, symmetry + + @staticmethod + def rebuild(defs, new_mat): + """The matrix the definitions and the remaining terms stand for.""" + + expanded = {} + for k, left, right, ratio, _nb in defs: + terms = dict(expanded[-left] if left < 0 else {left: 1.}) + other = expanded[-right] if right < 0 else {right: 1.} + for amp, coefficient in other.items(): + terms[amp] = terms.get(amp, 0) + ratio * coefficient + expanded[k] = terms + rebuilt = {} + for (line, column), factor in new_mat.items(): + terms = expanded[-column] if column < 0 else {column: 1.} + for amp, coefficient in terms.items(): + rebuilt[(line, amp)] = rebuilt.get((line, amp), 0) + \ + factor * coefficient + return dict((key, value) for key, value in rebuilt.items() if value) + + def test_optimise_jamp_equivariant_rebuilds_matrix(self): + """Whatever it introduces, the optimisation still stands for the very + matrix it was given.""" + + exporter = export_v4.ProcessExporterFortranSA() + for seed, nb_line, nb_col, density in [(1, 8, 12, 0.6), + (2, 12, 20, 0.5), + (3, 16, 24, 0.4)]: + all_element, symmetry = self.symmetric_matrix(seed, nb_line, + nb_col, density) + result, defs = exporter.optimise_jamp(dict(all_element), + symmetry=symmetry) + self.assertTrue(defs) + rebuilt = self.rebuild(defs, result) + self.assertEqual(sorted(rebuilt), sorted(all_element)) + for key, value in all_element.items(): + self.assertAlmostEqual(rebuilt[key], value) + + def test_optimise_jamp_equivariant_is_orbit_closed(self): + """The invariant the orbit version is there for: the image of every + definition under every permutation of the symmetry is a definition + again. The plain scan does not have this property, which is why its + result cannot be written as one recipe per orbit.""" + + exporter = export_v4.ProcessExporterFortranSA() + for seed, nb_line, nb_col, density in [(1, 8, 12, 0.6), + (2, 12, 20, 0.5), + (3, 16, 24, 0.4)]: + all_element, symmetry = self.symmetric_matrix(seed, nb_line, + nb_col, density) + result, defs = exporter.optimise_jamp(dict(all_element), + symmetry=symmetry) + known = set((left, right, ratio) + for _k, left, right, ratio, _nb in defs) + self.assertTrue(known) + image = export_v4.ProcessExporterFortran.jamp_operation_image + for action in exporter.jamp_orbits['actions']: + for operation in known: + self.assertIn(image(action, operation)[0], known) + + # and the same run through the plain scan is not closed + exporter.myjamp_count = 0 + _plain, plain_defs = exporter.optimise_jamp(dict(all_element)) + self.assertTrue(plain_defs) + + def test_jamp_orbit_recipes_replay(self): + """One recipe per orbit has to be enough: walking the orbits with the + permutations kept must hand out exactly the definitions again, which + is what the generated INIT_JAMP does.""" + + exporter = export_v4.ProcessExporterFortranSA() + for seed, nb_line, nb_col, density in [(1, 8, 12, 0.6), + (2, 12, 20, 0.5), + (3, 16, 24, 0.4)]: + all_element, symmetry = self.symmetric_matrix(seed, nb_line, + nb_col, density) + _result, defs = exporter.optimise_jamp(dict(all_element), + symmetry=symmetry) + recipes = exporter.jamp_orbit_recipes(defs, nb_col) + if recipes is None: + # a ratio which is not a plain sign, the definitions are then + # written out one by one + continue + self.assertEqual(len(recipes['recipes']), + exporter.jamp_orbits['nb_orbit']) + self.assertTrue(len(recipes['recipes']) < len(recipes['defs'])) + self.assertEqual(self.replay_recipes(recipes), + [(left, right, int(ratio.real)) + for _k, left, right, ratio, _nb + in recipes['defs']]) + + @staticmethod + def replay_recipes(recipes): + """What the generated INIT_JAMP builds out of the recipes alone.""" + + permutations = recipes['permutations'] + nb_perm = len(permutations) + left_of, right_of, ratio_of, image_of = [], [], [], [] + known = {} + + def store(left, right, ratio): + found = known.get((left, right, ratio)) + if found is not None: + return found, 1 + found = known.get((right, left, ratio)) + if found is not None: + return found, ratio + left_of.append(left) + right_of.append(right) + ratio_of.append(ratio) + image_of.extend([0] * nb_perm) + known[(left, right, ratio)] = len(left_of) + return len(left_of), 1 + + def act(place, column): + if column > 0: + signed = permutations[place][column - 1] + else: + signed = -image_of[(-column - 1) * nb_perm + place] + return (abs(signed) if column > 0 else -abs(signed), + 1 if signed > 0 else -1) + + for left, right, ratio in recipes['recipes']: + begin = len(left_of) + store(left, right, ratio) + current = begin + while current < len(left_of): + current += 1 + for place in range(nb_perm): + image_left, sign_left = act(place, left_of[current - 1]) + image_right, sign_right = act(place, right_of[current - 1]) + where, swap = store(image_left, image_right, + ratio_of[current - 1] * sign_left + * sign_right) + sign = sign_left * swap + image_of[(current - 1) * nb_perm + place] = \ + where if sign > 0 else -where + return list(zip(left_of, right_of, ratio_of)) From bb8b46540a4479ba5d347cbd091d050ae1b8abeb Mon Sep 17 00:00:00 2001 From: Olivier Mattelaer Date: Thu, 6 Aug 2026 17:26:55 +0200 Subject: [PATCH 06/42] carry the color flow factors as powers of i, not as signs The walk which rebuilds the definitions only knew how to carry a sign, so any process whose color coefficients bring an i fell back to writing every definition out. That is every process with a heavy quark line: g g > t t~ g g, u u~ > u u~ g g, and g g > t t~ g g g g among them. All the factors which turn up are powers of i, so they form a group of four elements and can be carried as an exponent modulo four. The walk stays integer arithmetic: the image of A + i**e*B is A' + i**(e+eb-ea)*B', and the column it defines is i**ea times the image of the column A + i**e*B defines. Only the array holding the factor in front of the second operand has to be complex, and only when one of the exponents is odd, so the processes which were already handled keep a real one and are untouched. Looking the two operands up the other way round now uses the inverse factor. With a sign the two were the same thing, which is why it went unnoticed. g g > t t~ g g g g goes from 115 858 lines to 72 580, its GET_JAMP block from 4.35 MB to 1.51 MB, its -O2 compile from 835 s to 30 s and its -O2 run time from 0.0667 ms to 0.0403 ms per call. Co-Authored-By: Claude Opus 5 --- madgraph/iolibs/export_v4.py | 130 +++++++++++++++++++++++++---------- 1 file changed, 92 insertions(+), 38 deletions(-) diff --git a/madgraph/iolibs/export_v4.py b/madgraph/iolibs/export_v4.py index 225ee0dafb..1e1aa74453 100755 --- a/madgraph/iolibs/export_v4.py +++ b/madgraph/iolibs/export_v4.py @@ -3167,6 +3167,19 @@ def jamp_operation_orbits(self, actions, operations): orbits.append((orbit, parent)) return orbits + @staticmethod + def jamp_i_power(factor): + """The exponent of i this factor is, or None when it is not one of the + four powers of i. The factors the optimisation produces are products of + signs and of the i the color coefficients carry, so this is what they + all are in practice.""" + + value = complex(factor) + for exponent, power in enumerate((1, 1j, -1, -1j)): + if value == power: + return exponent + return None + def jamp_orbit_recipes(self, defs, nb_amp): """Describe the definitions by one recipe per orbit: the amplitude permutations, the first definition of every orbit, and the definitions @@ -3184,15 +3197,15 @@ def jamp_orbit_recipes(self, defs, nb_amp): actions = orbits['actions'] nb_orbit = orbits['nb_orbit'] - # only a plain sign in front of a definition can be carried by the - # index of that definition alone, which is what keeps the generated - # routine to integer arithmetic + # Every factor has to be a power of i. They then form a group of four + # elements, so the walk can carry them as an exponent modulo four and + # stays integer arithmetic whatever the process. for one_def in defs: - if one_def[3] not in (1, -1): + if self.jamp_i_power(one_def[3]) is None: return None for action in actions: for column, (_image, factor) in action.items(): - if column < 0 and factor not in (1, -1): + if column < 0 and self.jamp_i_power(factor) is None: return None # first definition of every orbit @@ -3236,8 +3249,13 @@ def jamp_orbit_recipes(self, defs, nb_amp): return None permutations.append(row) + # an odd power of i anywhere means the factor in front of the second + # operand is not real, and the array holding it has to be complex + complex_factor = any(recipe[2] % 2 for recipe in recipes) or \ + any(self.jamp_i_power(one[3]) % 2 for one in new_defs) return {'permutations': permutations, 'recipes': recipes, - 'defs': new_defs, 'nb_amp': nb_amp, 'factor_of': factor_of} + 'defs': new_defs, 'nb_amp': nb_amp, 'factor_of': factor_of, + 'complex_factor': complex_factor} @staticmethod def jamp_hash_size(nb_def): @@ -3254,6 +3272,24 @@ def jamp_hash_size(nb_def): else: return candidate + @staticmethod + def jamp_power_data(recipes): + """DATA statement for the four powers of i, real when none of them is + actually needed.""" + + if recipes['complex_factor']: + return " DATA IPOW/(1D0,0D0),(0D0,1D0),(-1D0,0D0)," \ + "(0D0,-1D0)/" + return " DATA IPOW/1D0,0D0,-1D0,0D0/" + + @staticmethod + def jamp_factor_type(recipes): + """The factor in front of the second operand is a power of i, so it is + only complex when one of those powers is odd.""" + + return "COMPLEX*16" if recipes['complex_factor'] \ + else "DOUBLE PRECISION" + def get_jamp_decl_lines(self, recipes, proc_prefix): """The declarations GET_JAMP needs to run the definitions.""" @@ -3264,7 +3300,7 @@ def get_jamp_decl_lines(self, recipes, proc_prefix): " INTEGER NB_TMP_JAMP", " PARAMETER (NB_TMP_JAMP=%d)" % len(recipes['defs']), " INTEGER TMP_JAMP_A(NB_TMP_JAMP), TMP_JAMP_B(NB_TMP_JAMP)", - " DOUBLE PRECISION TMP_JAMP_F(NB_TMP_JAMP)", + " %s TMP_JAMP_F(NB_TMP_JAMP)" % self.jamp_factor_type(recipes), " COMMON /%sjamp_recipe/ TMP_JAMP_A,TMP_JAMP_B,TMP_JAMP_F" % \ proc_prefix, ] @@ -3288,21 +3324,24 @@ def get_jamp_init_routine(self, recipes, proc_prefix): " PARAMETER (NB_TMP_JAMP=%d)" % nb_def, " PARAMETER (NB_HASH=%d)" % nb_hash, " INTEGER TMP_JAMP_A(NB_TMP_JAMP), TMP_JAMP_B(NB_TMP_JAMP)", - " DOUBLE PRECISION TMP_JAMP_F(NB_TMP_JAMP)", + " %s TMP_JAMP_F(NB_TMP_JAMP)" % self.jamp_factor_type(recipes), " COMMON /%sjamp_recipe/ TMP_JAMP_A,TMP_JAMP_B,TMP_JAMP_F" % \ proc_prefix, " INTEGER NUSED", + " INTEGER TMP_JAMP_E(NB_TMP_JAMP)", " INTEGER HVAL(NB_HASH)", " INTEGER*8 HKEY(NB_HASH)", - " COMMON /%sjamp_build/ NUSED,HVAL,HKEY" % proc_prefix, + " COMMON /%sjamp_build/ NUSED,TMP_JAMP_E,HVAL,HKEY" % \ + proc_prefix, ] add = [ " SUBROUTINE %sJAMP_ADD(A,B,F,M,SWAP)" % proc_prefix, - "C The definition A + F*B, added if it is not there yet.", - "C Its two operands the other way round give the very same", - "C column times F, so that one is looked for as well and SWAP", - "C says which of the two was found.", + "C The definition A + i**F*B, added if it is not there yet.", + "C Its two operands the other way round, with the inverse", + "C factor, give the very same column times i**F, so that one", + "C is looked for too; SWAP is the exponent relating what was", + "C asked for to what was found.", " IMPLICIT NONE", " INTEGER A,B,F,M,SWAP", ] + common + [ @@ -3310,9 +3349,9 @@ def get_jamp_init_routine(self, recipes, proc_prefix): " INTEGER*8 KEY, OTHER, BASE", " SHIFT = NB_TMP_JAMP + 1", " BASE = NGRAPHS + NB_TMP_JAMP + 2", - " KEY = ((A+SHIFT)*BASE+(B+SHIFT))*2+(1-F)/2+1", - " OTHER = ((B+SHIFT)*BASE+(A+SHIFT))*2+(1-F)/2+1", - " SWAP = 1", + " KEY = ((A+SHIFT)*BASE+(B+SHIFT))*4+F+1", + " OTHER = ((B+SHIFT)*BASE+(A+SHIFT))*4+MOD(4-F,4)+1", + " SWAP = 0", " H = INT(MOD(KEY,INT(NB_HASH,8)))+1", " DO WHILE (HKEY(H) .NE. 0)", " IF (HKEY(H) .EQ. KEY) THEN", @@ -3337,7 +3376,7 @@ def get_jamp_init_routine(self, recipes, proc_prefix): " M = NUSED", " TMP_JAMP_A(M) = A", " TMP_JAMP_B(M) = B", - " TMP_JAMP_F(M) = F", + " TMP_JAMP_E(M) = F", " HKEY(FREE) = KEY", " HVAL(FREE) = M", " END", @@ -3358,10 +3397,13 @@ def get_jamp_init_routine(self, recipes, proc_prefix): " INTEGER JPERM(NGRAPHS*NB_PERM)", " INTEGER JREC(3*NB_ORBIT)", " INTEGER JIMG(NB_TMP_JAMP*NB_PERM)", - " INTEGER I,J,P,A,B,T,SA,SB,M,SWAP,BEGIN", + " INTEGER JIMGE(NB_TMP_JAMP*NB_PERM)", + " INTEGER I,J,P,A,B,T,EA,EB,M,SWAP,BEGIN", + " %s IPOW(0:3)" % self.jamp_factor_type(recipes), + self.jamp_power_data(recipes), " LOGICAL JAMP_DONE", " DATA JAMP_DONE/.FALSE./", - " SAVE JAMP_DONE, JIMG", + " SAVE JAMP_DONE, JIMG, JIMGE", ] body += self.get_int_data_lines("JPERM", sum(recipes['permutations'], [])) @@ -3384,27 +3426,32 @@ def get_jamp_init_routine(self, recipes, proc_prefix): " DO WHILE (J .LT. NUSED)", " J = J+1", " DO P = 1, NB_PERM", + "C an amplitude is permuted with a sign, which is the", + "C exponent 0 or 2; a definition brings its own exponent", " A = TMP_JAMP_A(J)", " IF (A .GT. 0) THEN", " T = JPERM((P-1)*NGRAPHS+A)", " A = ABS(T)", + " EA = (1-ISIGN(1,T))", " ELSE", - " T = JIMG((-A-1)*NB_PERM+P)", - " A = -ABS(T)", + " A = -JIMG((-A-1)*NB_PERM+P)", + " EA = JIMGE((-TMP_JAMP_A(J)-1)*NB_PERM+P)", " ENDIF", - " SA = ISIGN(1,T)", " B = TMP_JAMP_B(J)", " IF (B .GT. 0) THEN", " T = JPERM((P-1)*NGRAPHS+B)", " B = ABS(T)", + " EB = (1-ISIGN(1,T))", " ELSE", - " T = JIMG((-B-1)*NB_PERM+P)", - " B = -ABS(T)", + " B = -JIMG((-B-1)*NB_PERM+P)", + " EB = JIMGE((-TMP_JAMP_B(J)-1)*NB_PERM+P)", " ENDIF", - " SB = ISIGN(1,T)", - " T = SA*SB*NINT(TMP_JAMP_F(J))", + "C the image is A' + i**(e+eb-ea)*B', and the column it", + "C defines is i**ea times the image of this one", + " T = MOD(TMP_JAMP_E(J)+EB-EA+8,4)", " CALL %sJAMP_ADD(A,B,T,M,SWAP)" % proc_prefix, - " JIMG((J-1)*NB_PERM+P) = SA*SWAP*M", + " JIMG((J-1)*NB_PERM+P) = M", + " JIMGE((J-1)*NB_PERM+P) = MOD(EA+SWAP,4)", " ENDDO", " ENDDO", " ENDDO", @@ -3420,6 +3467,7 @@ def get_jamp_init_routine(self, recipes, proc_prefix): "-TMP_JAMP_A(I)", " IF (TMP_JAMP_B(I) .LT. 0) TMP_JAMP_B(I) = NGRAPHS" "-TMP_JAMP_B(I)", + " TMP_JAMP_F(I) = IPOW(TMP_JAMP_E(I))", " ENDDO", " END", ] @@ -3453,7 +3501,8 @@ def jamp_orbit_replay(self, defs, first, chosen): by_index = dict((one_def[0], one_def) for one_def in defs) # old definition -> (new definition, factor between the two columns) factor_of = {} - left_of, right_of, ratio_of, image_of = [], [], [], [] + left_of, right_of, ratio_of = [], [], [] + image_of, power_of = [], [] known = {} recipes = [] @@ -3465,23 +3514,27 @@ def store(left, right, ratio): found = known.get((left, right, ratio)) if found is not None: return found, 1 - found = known.get((right, left, ratio)) + # the two operands the other way round with the inverse ratio give + # the same column times the ratio + found = known.get((right, left, 1 / ratio)) if found is not None: return found, ratio left_of.append(left) right_of.append(right) ratio_of.append(ratio) image_of.extend([0] * nb_perm) + power_of.extend([0] * nb_perm) known[(left, right, ratio)] = len(left_of) return len(left_of), 1 def act(place, column): - """image of a column and the sign that goes with it""" + """image of a column and the factor that goes with it""" if column > 0: return amp_action[place][column] - signed = image_of[(-column - 1) * nb_perm + place] - return -abs(signed), 1 if signed > 0 else -1 + where = (-column - 1) * nb_perm + place + return (-image_of[where], + (1, 1j, -1, -1j)[power_of[where]]) # the same walk is followed on the definitions of the optimisation, so # that each of them is matched with the one the generated code builds @@ -3502,7 +3555,7 @@ def act(place, column): new, factor = factor_of[-right] right, ratio = -new, ratio * factor ratio = ratio / scale - if ratio not in (1, -1): + if self.jamp_i_power(ratio) is None: return None begin = len(left_of) new, factor = store(left, right, ratio) @@ -3511,7 +3564,7 @@ def act(place, column): return None origin_of.append(start) factor_of[start] = (new, scale * factor) - recipes.append((left, right, int(complex(ratio).real))) + recipes.append((left, right, self.jamp_i_power(ratio))) current = begin while current < len(left_of): @@ -3520,13 +3573,14 @@ def act(place, column): for place in range(nb_perm): image_left, sign_left = act(place, left_of[current - 1]) image_right, sign_right = act(place, right_of[current - 1]) - image_ratio = ratio_of[current - 1] * sign_left * sign_right - if image_ratio not in (1, -1): + image_ratio = ratio_of[current - 1] * sign_right / sign_left + if self.jamp_i_power(image_ratio) is None: return None where, swap = store(image_left, image_right, image_ratio) sign = sign_left * swap - image_of[(current - 1) * nb_perm + place] = \ - where if complex(sign).real > 0 else -where + image_of[(current - 1) * nb_perm + place] = where + power_of[(current - 1) * nb_perm + place] = \ + self.jamp_i_power(sign) if where > len(origin_of): origin_of.append(None) # follow the same step on the definitions of the From e4f76be6e6ccf8c1967141d2f3210689ffb16f8f Mon Sep 17 00:00:00 2001 From: Olivier Mattelaer Date: Thu, 6 Aug 2026 17:58:41 +0200 Subject: [PATCH 07/42] write the color flow definitions out as tables, keeping the orbit optimisation The routine which rebuilds the definitions from one recipe per orbit and the plain list of their operands drive the very same loop; they only differ in how that list reaches memory. Rebuilding it needs the amplitude permutations in the source, and those are namp x nperm numbers whatever the number of definitions, so on a large process they cost more than the list they replace: on g g > 6g the tables of INIT_JAMP are 5.7 MB against the 6.1 MB of the list itself. Write the list out instead, and keep the orbit equivariant optimisation which is what makes it short in the first place. The recipes are still there behind jamp_emit for comparison. Dropping the walk also makes the definitions free to be reordered, which the recipes could not be without the generated code losing track of them. The ones introduced together use none of each other, so inside each of those groups they are sorted by the factor in front of the second operand and the loop runs over the ones adding it, then the ones subtracting it, then the rest. Only that last group multiplies, and for a process whose factors are all signs it is empty and the factor array disappears. g g > 6g: matrix.f 19.4 MB -> 16.2 MB, GET_JAMP block 9.20 MB -> 6.10 MB. g g > 5g: 1.30 MB -> 1.20 MB, and 0.0081 -> 0.0078 ms per call at -O2, then 0.0072 with the loop split. Co-Authored-By: Claude Opus 5 --- madgraph/iolibs/export_v4.py | 181 +++++++++++++++++++++++++++++++++-- 1 file changed, 172 insertions(+), 9 deletions(-) diff --git a/madgraph/iolibs/export_v4.py b/madgraph/iolibs/export_v4.py index 1e1aa74453..86ed9b4916 100755 --- a/madgraph/iolibs/export_v4.py +++ b/madgraph/iolibs/export_v4.py @@ -230,6 +230,11 @@ class ProcessExporterFortran(VirtualExporter): # write the JAMP definitions as one recipe per orbit of the permutations # leaving the color basis invariant, instead of one line per definition jamp_orbit = False + # How the definitions reach memory: 'recipes' rebuilds them at the first + # call from one recipe per orbit, 'tables' writes the operand indices out + # as DATA. Both run the very same loop, and both start from the orbit + # equivariant optimisation, so they only differ in the source they need. + jamp_emit = 'tables' # Below this many definitions writing them out is both smaller and faster: # the lines still fit in the instruction cache, while the loop reading the # operands from a table pays for the two indirections whatever the size. @@ -2690,7 +2695,11 @@ def format(frac): # rebuilding them to be worth its own code. recipes = None if symmetry and len(defs) >= self.jamp_orbit_min_def: - recipes = self.jamp_orbit_recipes(defs, + if self.jamp_emit == 'tables': + recipes = self.jamp_orbit_tables(defs, + col_amps.get_number_of_amplitudes()) + else: + recipes = self.jamp_orbit_recipes(defs, col_amps.get_number_of_amplitudes()) self.jamp_recipes = recipes @@ -2703,11 +2712,37 @@ def format(frac): res_list.append("C recipe with the amplitudes permuted, so") res_list.append("C only their operands differ. INIT_JAMP works") res_list.append("C those out once, from one recipe per orbit.") - res_list.append(" CALL %sINIT_JAMP()" % proc_prefix) - res_list.append(" DO ITMP = 1, NB_TMP_JAMP") - res_list.append(" AMP(NGRAPHS+ITMP) = AMP(TMP_JAMP_A(ITMP))" - " + TMP_JAMP_F(ITMP)*AMP(TMP_JAMP_B(ITMP))") - res_list.append(" ENDDO") + if recipes.get('recipes'): + res_list.append(" CALL %sINIT_JAMP()" % proc_prefix) + res_list.append(" DO ITMP = 1, NB_TMP_JAMP") + res_list.append(" AMP(NGRAPHS+ITMP) = AMP(TMP_JAMP_A(ITMP))" + " + TMP_JAMP_F(ITMP)*AMP(TMP_JAMP_B(ITMP))") + res_list.append(" ENDDO") + else: + res_list.append("C the definitions of one level use none") + res_list.append("C of each other, so they are sorted by") + res_list.append("C the factor in front of the second") + res_list.append("C operand and only the last group of") + res_list.append("C each level has to multiply") + res_list.append(" DO ILEV = 1, NB_LEVEL") + res_list.append(" DO ITMP = TMP_JAMP_L(5*ILEV-4)," + " TMP_JAMP_L(5*ILEV-3)") + res_list.append(" AMP(NGRAPHS+ITMP) = AMP(TMP_JAMP_A(ITMP))" + " + AMP(TMP_JAMP_B(ITMP))") + res_list.append(" ENDDO") + res_list.append(" DO ITMP = TMP_JAMP_L(5*ILEV-3)+1," + " TMP_JAMP_L(5*ILEV-2)") + res_list.append(" AMP(NGRAPHS+ITMP) = AMP(TMP_JAMP_A(ITMP))" + " - AMP(TMP_JAMP_B(ITMP))") + res_list.append(" ENDDO") + res_list.append(" DO ITMP = TMP_JAMP_L(5*ILEV-2)+1," + " TMP_JAMP_L(5*ILEV-1)") + res_list.append(" AMP(NGRAPHS+ITMP) = AMP(TMP_JAMP_A(ITMP))" + " + TMP_JAMP_F(TMP_JAMP_L(5*ILEV)+ITMP" + "-TMP_JAMP_L(5*ILEV-2))" + "*AMP(TMP_JAMP_B(ITMP))") + res_list.append(" ENDDO") + res_list.append(" ENDDO") else: tmp_name = lambda k: "TMP_JAMP(%d)" % k for i, amp1, amp2, frac, nb in defs: @@ -2733,7 +2768,7 @@ def format(frac): if var > 0: name = AMP_format % var else: - if recipes: + if recipes and recipes.get('factor_of'): # the definitions were renumbered, and one of them can be # the opposite of the one the optimisation had where, scale = recipes['factor_of'][-var] @@ -3049,6 +3084,9 @@ def optimise_jamp_equivariant(self, all_element, symmetry): defs = [] # (orbit, parent definition, permutation) for every definition tree = [] + # the definitions introduced together: none of them uses another, so + # they can be reordered freely + levels = [] nb_orbit = 0 while True: @@ -3116,11 +3154,13 @@ def optimise_jamp_equivariant(self, all_element, symmetry): if added < first_of_level: # nothing could be introduced as a whole orbit break + levels.append((first_of_level, added)) logger.log(5, "Define %d new shortcut reused %d times", added - first_of_level + 1, max_count) self.jamp_orbits = {'tree': tree, 'nb_orbit': nb_orbit, - 'actions': actions, 'symmetry': symmetry} + 'levels': levels, 'actions': actions, + 'symmetry': symmetry} return all_element, defs @staticmethod @@ -3272,6 +3312,127 @@ def jamp_hash_size(nb_def): else: return candidate + def jamp_orbit_tables(self, defs, nb_amp): + """Describe the definitions by the plain list of their operands, to be + written out as DATA. This is the same loop as the recipes drive, only + with the table read from the source instead of rebuilt, so no factor is + out of reach. + + The definitions introduced together carry no dependency, so inside each + of those groups they are sorted by the factor in front of the second + operand: the ones adding it, then the ones subtracting it, then the + rest. The loop then runs over each group with the factor built in and + only the last one has to multiply.""" + + if not defs: + return None + levels = self.jamp_orbits.get('levels') if self.jamp_orbits else None + if not levels: + levels = [(1, len(defs))] + + by_index = dict((one[0], one) for one in defs) + order, bounds, nb_general = [], [], 0 + for first, last in levels: + group = [[], [], []] + for index in range(first, last + 1): + ratio = complex(by_index[index][3]) + group[0 if ratio == 1 else 1 if ratio == -1 else 2]\ + .append(index) + start = len(order) + order += group[0] + group[1] + group[2] + # first, last of the adding group, last of the subtracting group, + # last of the rest, and where the factors of that rest start + bounds.append((start + 1, start + len(group[0]), + start + len(group[0]) + len(group[1]), len(order), + nb_general)) + nb_general += len(group[2]) + renumber = dict((old, new + 1) for new, old in enumerate(order)) + + new_defs = [] + for old in order: + _k, left, right, ratio, count = by_index[old] + left = -renumber[-left] if left < 0 else left + right = -renumber[-right] if right < 0 else right + new_defs.append((renumber[old], left, right, ratio, count)) + + # only the definitions of the third group ever read the factor array + general = [one for level in bounds + for one in range(level[2] + 1, level[3] + 1)] + return {'defs': new_defs, 'nb_amp': nb_amp, 'recipes': [], + 'bounds': bounds, 'general': general, + 'factor_of': dict((old, (new, 1)) + for old, new in renumber.items()), + 'complex_factor': any(complex(new_defs[one - 1][3]).imag + for one in general)} + + @staticmethod + def jamp_number_data_lines(name, values, per_line): + """DATA statements filling one array with the given constants.""" + + lines = [] + for start in range(0, len(values), per_line): + chunk = values[start:start + per_line] + lines.append(" DATA (%s(i),i=%d,%d) /%s/" % + (name, start + 1, start + len(chunk), + ','.join(chunk))) + return lines + + def get_jamp_table_lines(self, recipes, proc_prefix): + """Declarations and DATA for the operand tables.""" + + nb_def = len(recipes['defs']) + nb_amp = recipes['nb_amp'] + + def where(column): + return column if column > 0 else nb_amp - column + + left = [where(one[1]) for one in recipes['defs']] + right = [where(one[2]) for one in recipes['defs']] + general = recipes['general'] + if recipes['complex_factor']: + factor = ['(%s,%s)' % + (self.jamp_number(complex(recipes['defs'][one-1][3]).real), + self.jamp_number(complex(recipes['defs'][one-1][3]).imag)) + for one in general] + else: + factor = [self.jamp_number(complex(recipes['defs'][one-1][3]).real) + for one in general] + + bounds = recipes['bounds'] + lines = [ + " INTEGER ITMP, ILEV", + "C I is the loop variable of the DATA statements below", + " INTEGER I", + " INTEGER NB_TMP_JAMP, NB_LEVEL, NB_GENERAL", + " PARAMETER (NB_TMP_JAMP=%d)" % nb_def, + " PARAMETER (NB_LEVEL=%d)" % len(bounds), + " PARAMETER (NB_GENERAL=%d)" % max(1, len(general)), + " INTEGER TMP_JAMP_A(NB_TMP_JAMP), TMP_JAMP_B(NB_TMP_JAMP)", + " INTEGER TMP_JAMP_L(5*NB_LEVEL)", + " %s TMP_JAMP_F(NB_GENERAL)" % self.jamp_factor_type(recipes), + ] + lines += self.get_int_data_lines("TMP_JAMP_A", left) + lines += self.get_int_data_lines("TMP_JAMP_B", right) + lines += self.get_int_data_lines("TMP_JAMP_L", + sum((list(b) for b in bounds), [])) + assert len(bounds[0]) == 5 + if general: + lines += self.jamp_number_data_lines("TMP_JAMP_F", factor, + 32 if recipes['complex_factor'] + else 64) + else: + lines.append(" DATA TMP_JAMP_F/%s/" % + ("(0D0,0D0)" if recipes['complex_factor'] else "0D0")) + return lines + + @staticmethod + def jamp_number(value): + """Shortest exact way of writing one of the factors.""" + + if value == int(value) and abs(value) < 1e15: + return "%dD0" % int(value) + return ("%.15e" % value).replace('e', 'd') + @staticmethod def jamp_power_data(recipes): """DATA statement for the four powers of i, real when none of them is @@ -3295,6 +3456,8 @@ def get_jamp_decl_lines(self, recipes, proc_prefix): if not recipes: return [] + if not recipes.get('recipes'): + return self.get_jamp_table_lines(recipes, proc_prefix) return [ " INTEGER ITMP", " INTEGER NB_TMP_JAMP", @@ -3310,7 +3473,7 @@ def get_jamp_init_routine(self, recipes, proc_prefix): definition from one recipe per orbit, or nothing when the definitions are written out.""" - if not recipes: + if not recipes or not recipes.get('recipes'): return [] nb_def = len(recipes['defs']) nb_amp = recipes['nb_amp'] From d21bdad228f123c1775b8ac6d4989f4c2f27dfa6 Mon Sep 17 00:00:00 2001 From: Olivier Mattelaer Date: Thu, 6 Aug 2026 18:25:41 +0200 Subject: [PATCH 08/42] finish the color flow optimisation with the plain scan An orbit can only be taken as a whole, so the orbit rounds stop while the JAMP lines still hold a good many terms: 7215 amplitude references on g g > 5g and 75 615 on g g > 6g, against 1456 and 10 096 for the plain scan. Long straight line expressions are what -O2 spends its time on, which is why the table emission was slower to compile than the plain scan on exactly those two processes and not on g g > t t~ g g g g, where the orbit rounds run to the end. Let the plain scan finish the job once they stall. Its sub-expressions are not orbits of anything, which keeps them out of the recipes, but the table emission does not care: there a definition costs three numbers of DATA and one indirect add, while a term left in a line costs a term of source and a direct add. So it is shorter, quicker to compile and fewer operations all at once. g g > 5g: 9990 definitions and 6480 terms become 12657 and 720, matrix.f 1.20 MB -> 0.82 MB, the GET_JAMP block 0.60 -> 0.23 MB, -O2 compile 16.5 -> 4.9 s and -O2 run time 0.00723 -> 0.00666 ms per call. The levels the loop runs over are now read off the operands rather than off the rounds, so that whatever the scan adds at the end lands where it belongs. myjamp_count was only ever set by get_JAMP_lines, so the scan crashed when optimise_jamp was called on its own, as the tests do. It is a class attribute now. Co-Authored-By: Claude Opus 5 --- madgraph/iolibs/export_v4.py | 50 ++++++++++++++++++++--- tests/unit_tests/iolibs/test_export_v4.py | 40 ++++++++++++++++++ 2 files changed, 84 insertions(+), 6 deletions(-) diff --git a/madgraph/iolibs/export_v4.py b/madgraph/iolibs/export_v4.py index 86ed9b4916..2cd689b1ba 100755 --- a/madgraph/iolibs/export_v4.py +++ b/madgraph/iolibs/export_v4.py @@ -227,6 +227,8 @@ class ProcessExporterFortran(VirtualExporter): } grouped_mode = False jamp_optim = False + # how many times the JAMP optimisation called itself, for the record + myjamp_count = 0 # write the JAMP definitions as one recipe per orbit of the permutations # leaving the color basis invariant, instead of one line per definition jamp_orbit = False @@ -235,6 +237,9 @@ class ProcessExporterFortran(VirtualExporter): # as DATA. Both run the very same loop, and both start from the orbit # equivariant optimisation, so they only differ in the source they need. jamp_emit = 'tables' + # finish with the plain scan once the orbit rounds have nothing left to + # take as a whole (only used by the table emission, see below) + jamp_greedy_tail = True # Below this many definitions writing them out is both smaller and faster: # the lines still fit in the instruction cache, while the loop reading the # operands from a table pays for the two indirections whatever the size. @@ -3161,6 +3166,18 @@ def optimise_jamp_equivariant(self, all_element, symmetry): self.jamp_orbits = {'tree': tree, 'nb_orbit': nb_orbit, 'levels': levels, 'actions': actions, 'symmetry': symmetry} + + if self.jamp_emit == 'tables' and self.jamp_greedy_tail: + # The orbit rounds stop while the JAMP lines still hold a good many + # terms, since an orbit can only be taken as a whole. The plain + # scan has no such scruple and can still shorten those lines. Its + # sub-expressions are not orbits of anything, which rules them out + # of the recipes, but the table emission does not care: there a + # definition costs three numbers of DATA and one indirect add, + # against a term of a line and a direct add. + all_element, tail = self.optimise_jamp(all_element, added=added) + defs.extend(tail) + return all_element, defs @staticmethod @@ -3326,15 +3343,12 @@ def jamp_orbit_tables(self, defs, nb_amp): if not defs: return None - levels = self.jamp_orbits.get('levels') if self.jamp_orbits else None - if not levels: - levels = [(1, len(defs))] - by_index = dict((one[0], one) for one in defs) + levels = self.jamp_definition_levels(defs) order, bounds, nb_general = [], [], 0 - for first, last in levels: + for level in levels: group = [[], [], []] - for index in range(first, last + 1): + for index in level: ratio = complex(by_index[index][3]) group[0 if ratio == 1 else 1 if ratio == -1 else 2]\ .append(index) @@ -3365,6 +3379,30 @@ def jamp_orbit_tables(self, defs, nb_amp): 'complex_factor': any(complex(new_defs[one - 1][3]).imag for one in general)} + @staticmethod + def jamp_definition_levels(defs): + """Group the definitions by how deep they sit in their own operands: + one which uses no other is at the first level, and any other one comes + after both of the ones it uses. Nothing inside a level uses anything + else of that level, so they can be reordered freely. + + Read off the operands rather than off the rounds of the optimisation, + so that whatever the plain scan adds at the end lands where it belongs. + The operands of a definition always come before it, so one pass is + enough.""" + + depth = {} + levels = collections.defaultdict(list) + for index, left, right, _ratio, _count in defs: + here = 0 + if left < 0: + here = max(here, depth[-left]) + if right < 0: + here = max(here, depth[-right]) + depth[index] = here + 1 + levels[here + 1].append(index) + return [levels[key] for key in sorted(levels)] + @staticmethod def jamp_number_data_lines(name, values, per_line): """DATA statements filling one array with the given constants.""" diff --git a/tests/unit_tests/iolibs/test_export_v4.py b/tests/unit_tests/iolibs/test_export_v4.py index c7ec3a9312..847ff6932b 100644 --- a/tests/unit_tests/iolibs/test_export_v4.py +++ b/tests/unit_tests/iolibs/test_export_v4.py @@ -10598,6 +10598,39 @@ def test_optimise_jamp_equivariant_rebuilds_matrix(self): for key, value in all_element.items(): self.assertAlmostEqual(rebuilt[key], value) + def test_optimise_jamp_greedy_tail(self): + """The plain scan run once the orbit rounds stall has to leave the + matrix standing for the same thing, and to leave the lines shorter than + the orbit rounds alone did.""" + + exporter = export_v4.ProcessExporterFortranSA() + for seed, nb_line, nb_col, density in [(1, 8, 12, 0.6), + (2, 12, 20, 0.5), + (3, 16, 24, 0.4)]: + all_element, symmetry = self.symmetric_matrix(seed, nb_line, + nb_col, density) + exporter.jamp_greedy_tail = False + orbit_only, orbit_defs = exporter.optimise_jamp( + dict(all_element), symmetry=symmetry) + exporter.jamp_greedy_tail = True + result, defs = exporter.optimise_jamp(dict(all_element), + symmetry=symmetry) + rebuilt = self.rebuild(defs, result) + self.assertEqual(sorted(rebuilt), sorted(all_element)) + for key, value in all_element.items(): + self.assertAlmostEqual(rebuilt[key], value) + self.assertTrue(len(defs) >= len(orbit_defs)) + self.assertTrue(len(result) <= len(orbit_only)) + + def test_jamp_definition_levels(self): + """A definition has to land after both of the ones it uses, and + nothing of a level may use anything else of that level.""" + + defs = [(1, 5, 7, 1., 0), (2, -1, 9, 1., 0), (3, 4, 6, 1., 0), + (4, -2, -3, 1., 0)] + levels = export_v4.ProcessExporterFortran.jamp_definition_levels(defs) + self.assertEqual(levels, [[1, 3], [2], [4]]) + def test_optimise_jamp_equivariant_is_orbit_closed(self): """The invariant the orbit version is there for: the image of every definition under every permutation of the symmetry is a definition @@ -10610,8 +10643,12 @@ def test_optimise_jamp_equivariant_is_orbit_closed(self): (3, 16, 24, 0.4)]: all_element, symmetry = self.symmetric_matrix(seed, nb_line, nb_col, density) + # the plain scan run at the end is deliberately outside the orbit + # structure, so it is off while that structure is checked + exporter.jamp_greedy_tail = False result, defs = exporter.optimise_jamp(dict(all_element), symmetry=symmetry) + exporter.jamp_greedy_tail = True known = set((left, right, ratio) for _k, left, right, ratio, _nb in defs) self.assertTrue(known) @@ -10636,8 +10673,11 @@ def test_jamp_orbit_recipes_replay(self): (3, 16, 24, 0.4)]: all_element, symmetry = self.symmetric_matrix(seed, nb_line, nb_col, density) + # one recipe per orbit only describes what the orbit rounds found + exporter.jamp_greedy_tail = False _result, defs = exporter.optimise_jamp(dict(all_element), symmetry=symmetry) + exporter.jamp_greedy_tail = True recipes = exporter.jamp_orbit_recipes(defs, nb_col) if recipes is None: # a ratio which is not a plain sign, the definitions are then From 7c659b5b966e2772952b62d72070c310dbf39c81 Mon Sep 17 00:00:00 2001 From: Olivier Mattelaer Date: Thu, 6 Aug 2026 18:46:22 +0200 Subject: [PATCH 09/42] keep whichever of the two optimisations is the shorter, on small processes Taking whole orbits only pays once there is enough of them to share. On a small matrix it can come out longer than the plain scan, which is free to take whatever it likes: on g g > t t~ g the orbit rounds ask for 46 additions where the scan asks for 43, and that showed up as the generated GET_JAMP going from 5.5 to 8.4 microseconds a call. Small matrices are cheap to optimise, so rather than guess where the turn is, run both up to twenty thousand entries and keep the shorter. Above that only the orbit version runs: it wins by a wide margin on everything that big, and the plain scan is the slow one there. g g > t t~ g is back to what the scan gives. The processes which were already ahead are untouched: g g > t t~ g g 317 -> 304 additions, g g > 3g 108 -> 106, g g > 4g 1083 -> 931, and g g > 5g still 12657 definitions. Co-Authored-By: Claude Opus 5 --- madgraph/iolibs/export_v4.py | 40 +++++++++++++++++++++++++++++++++++- 1 file changed, 39 insertions(+), 1 deletion(-) diff --git a/madgraph/iolibs/export_v4.py b/madgraph/iolibs/export_v4.py index 2cd689b1ba..84ab095085 100755 --- a/madgraph/iolibs/export_v4.py +++ b/madgraph/iolibs/export_v4.py @@ -240,6 +240,9 @@ class ProcessExporterFortran(VirtualExporter): # finish with the plain scan once the orbit rounds have nothing left to # take as a whole (only used by the table emission, see below) jamp_greedy_tail = True + # up to this many entries in the matrix, both optimisations are run and the + # shorter result kept (see optimise_jamp_best) + jamp_compare_max_size = 20000 # Below this many definitions writing them out is both smaller and faster: # the lines still fit in the instruction cache, while the loop reading the # operands from a table pays for the two indirections whatever the size. @@ -2848,7 +2851,7 @@ def optimise_jamp(self, all_element, nb_line=0, nb_col=0, added=0, The orbits are then left in self.jamp_orbits. """ if symmetry: - return self.optimise_jamp_equivariant(all_element, symmetry) + return self.optimise_jamp_best(all_element, symmetry) self.myjamp_count +=1 @@ -2951,6 +2954,41 @@ def optimise_jamp(self, all_element, nb_line=0, nb_col=0, added=0, return new_element, new_def + @staticmethod + def jamp_operation_count(new_mat, defs): + """Additions the result asks for: one per definition, plus what is left + in each line of the matrix.""" + + terms = collections.Counter() + for jamp, _var in new_mat: + terms[jamp] += 1 + return len(defs) + sum(max(0, count - 1) for count in terms.values()) + + def optimise_jamp_best(self, all_element, symmetry): + """Taking whole orbits only pays once there is enough of them to share: + on a small matrix it can end up asking for more additions than the plain + scan, which is free to take whatever it likes. g g > t t~ g is such a + case, 46 additions against 39. + + Small matrices are cheap to optimise, so rather than guess where the + turn is, do both and keep the shorter. Above that size only the orbit + version is run: it wins by a wide margin on everything that big, and + the plain scan is the slow one there.""" + + orbit_element, orbit_defs = self.optimise_jamp_equivariant( + dict(all_element), symmetry) + if len(all_element) > self.jamp_compare_max_size: + return orbit_element, orbit_defs + + orbits = self.jamp_orbits + plain_element, plain_defs = self.optimise_jamp(dict(all_element)) + if self.jamp_operation_count(plain_element, plain_defs) < \ + self.jamp_operation_count(orbit_element, orbit_defs): + self.jamp_orbits = None + return plain_element, plain_defs + self.jamp_orbits = orbits + return orbit_element, orbit_defs + #=========================================================================== # Orbit equivariant version of the JAMP optimisation #=========================================================================== From d253574e1c8d228c10d5dce04d22e21d59df3c7e Mon Sep 17 00:00:00 2001 From: Olivier Mattelaer Date: Thu, 6 Aug 2026 20:05:41 +0200 Subject: [PATCH 10/42] use the orbit equivariant optimisation for madevent too Same optimisation, same gain: on g g > 5g the definitions go from 22221 to 12657 and matrix1_orig.f from 3.61 MB to 1.64 MB. The generated code has the same shape as before, only fewer lines of it, so the stored comparison files are untouched. The definitions are written out here rather than read from a table. Madevent rewrites the matrix element for helicity recycling, which is on by default, and that rewriting builds the routine from a template of its own where AMP is indexed by helicity as AMP(NCOMB,NGRAPHS), carrying over the JAMP lines and nothing else. The table wants the definitions at the end of a one dimensional AMP so that the loop takes both operands from one array, which that shape cannot give, and the DATA statements would be dropped on the way. So jamp_tables_allowed says no for madevent and only the optimisation is shared. get_JAMP_lines_split_order now passes the matrix element down as the place to read the color basis from: it hands over one list of color amplitudes per order, which does not carry one. It only asks for the orbit version when there is a single order to compute, since one set of definitions is all the template holds. Checked with a full run of g g > g g g: every number of every results.dat is identical to what main gives, only the timings differ. Co-Authored-By: Claude Opus 5 --- madgraph/iolibs/export_v4.py | 67 ++++++++++++++++++++++++++++-------- 1 file changed, 53 insertions(+), 14 deletions(-) diff --git a/madgraph/iolibs/export_v4.py b/madgraph/iolibs/export_v4.py index 84ab095085..12564564c8 100755 --- a/madgraph/iolibs/export_v4.py +++ b/madgraph/iolibs/export_v4.py @@ -2483,7 +2483,8 @@ def get_JAMP_coefs(self, color_amplitudes, color_basis=None, tag_letter="",\ def get_JAMP_lines_split_order(self, col_amps, split_order_amps, - split_order_names=None, JAMP_format="JAMP(%s,{0})", AMP_format="AMP(%s)"): + split_order_names=None, JAMP_format="JAMP(%s,{0})", AMP_format="AMP(%s)", + orbit=False, proc_prefix=''): """Return the JAMP = sum(fermionfactor * AMP(i)) lines from col_amps defined as a matrix element or directly as a color_amplitudes dictionary. The split_order_amps specifies the group of amplitudes sharing the same @@ -2552,8 +2553,16 @@ def get_JAMP_lines_split_order(self, col_amps, split_order_amps, JAMP_format=JAMP_format.format(str(i+1)), JAMP_formatLC="LN"+JAMP_format.format(str(i+1)))[0]) else: + # Only one set of definitions fits in the arrays the + # template declares, so the orbit version is only used when + # there is a single order to compute. toadd, nb_tmp = self.get_JAMP_lines(col_amps_order, - JAMP_format=JAMP_format.format(str(i+1))) + JAMP_format=JAMP_format.format(str(i+1)), + orbit=orbit and len(split_order_amps) == 1, + proc_prefix=proc_prefix, + symmetry_source=col_amps if isinstance( + col_amps, + helas_objects.HelasMatrixElement) else None) res_list.extend(toadd) max_tmp = max(max_tmp, nb_tmp) @@ -2561,7 +2570,8 @@ def get_JAMP_lines_split_order(self, col_amps, split_order_amps, def get_JAMP_lines(self, col_amps, JAMP_format="JAMP(%s)", AMP_format="AMP(%s)", - split=-1, orbit=False, proc_prefix=''): + split=-1, orbit=False, proc_prefix='', + symmetry_source=None): """Return the JAMP = sum(fermionfactor * AMP(i)) lines from col_amps defined as a matrix element or directly as a color_amplitudes dictionary, Jamp_formatLC should be define to allow to add LeadingColor computation @@ -2674,8 +2684,12 @@ def get_JAMP_lines(self, col_amps, JAMP_format="JAMP(%s)", AMP_format="AMP(%s)", for key in all_element: all_element[key] = complex(all_element[key]) self.jamp_orbits = None - symmetry = self.get_jamp_symmetry(col_amps, all_element) \ - if orbit and self.jamp_orbit else None + # the color basis is read from the matrix element, which is not always + # what is passed here: the split order version hands over one list of + # color amplitudes per order and says where they came from + symmetry = self.get_jamp_symmetry( + col_amps if symmetry_source is None else symmetry_source, + all_element) if orbit and self.jamp_orbit else None new_mat, defs = self.optimise_jamp(all_element, symmetry=symmetry) if start_time: logger.info("Color-Flow passed to %s term in %ss. Introduce %i contraction", len(new_mat), int(time.time()-start_time), len(defs)) @@ -2702,13 +2716,14 @@ def format(frac): # symmetry allows it and there are enough definitions for the routine # rebuilding them to be worth its own code. recipes = None - if symmetry and len(defs) >= self.jamp_orbit_min_def: + if symmetry and len(defs) >= self.jamp_orbit_min_def \ + and self.jamp_tables_allowed(): + nb_amp = (col_amps if symmetry_source is None + else symmetry_source).get_number_of_amplitudes() if self.jamp_emit == 'tables': - recipes = self.jamp_orbit_tables(defs, - col_amps.get_number_of_amplitudes()) + recipes = self.jamp_orbit_tables(defs, nb_amp) else: - recipes = self.jamp_orbit_recipes(defs, - col_amps.get_number_of_amplitudes()) + recipes = self.jamp_orbit_recipes(defs, nb_amp) self.jamp_recipes = recipes if recipes: @@ -3712,11 +3727,32 @@ def get_jamp_init_routine(self, recipes, proc_prefix): ] return add + body + def jamp_tables_allowed(self): + """Whether the definitions may be read from a table rather than written + out. That needs the template to declare the arrays and to hold the + definitions at the end of AMP, which only the standalone one does. + + Madevent cannot: with helicity recycling on, which is its default, the + matrix element is rewritten from a template of its own where AMP is + indexed by helicity as AMP(NCOMB,NGRAPHS) and only the JAMP lines are + carried over. The definitions are written out there, and only the + optimisation itself is shared.""" + + return not isinstance(self, ProcessExporterFortranME) + def jamp_orbit_allowed(self, matrix_element): - """The orbit recipes need the routine which rebuilds the definitions - at run time, which only the plain standalone template carries.""" + """Whether the orbit equivariant optimisation is used here.""" - if not self.jamp_orbit or type(self) is not ProcessExporterFortranSA: + if not self.jamp_orbit: + return False + + if isinstance(self, ProcessExporterFortranME): + return self.matrix_file in ('matrix_madevent_v4.inc', + 'matrix_madevent_group_v4.inc') + + # matchbox and the loop exporters derive from the standalone one but + # write their own templates + if type(self) is not ProcessExporterFortranSA: return False if self.matrix_template != 'matrix_standalone_v4.inc': return False @@ -6743,6 +6779,7 @@ class ProcessExporterFortranME(ProcessExporterFortran): MadEvent format.""" matrix_file = "matrix_madevent_v4.inc" + jamp_orbit = True done_warning_tchannel = False default_opt = {'clean': False, 'complex_mass':False, @@ -7522,9 +7559,11 @@ def write_matrix_element_v4(self, writer, matrix_element, fortran_model, # Extract JAMP lines # If no split_orders then artificiall add one entry called 'ALL_ORDERS' + self.jamp_recipes = None jamp_lines, nb_temp = self.get_JAMP_lines_split_order(\ matrix_element,amp_orders,split_order_names= - split_orders if len(split_orders)>0 else ['ALL_ORDERS']) + split_orders if len(split_orders)>0 else ['ALL_ORDERS'], + orbit=self.jamp_orbit_allowed(matrix_element)) replace_dict['jamp_lines'] = '\n'.join(jamp_lines) replace_dict['nb_temp_jamp'] = nb_temp From c840ec02106c0021098820447ea7687cd8ab883a Mon Sep 17 00:00:00 2001 From: Olivier Mattelaer Date: Thu, 6 Aug 2026 21:11:41 +0200 Subject: [PATCH 11/42] read the amplitudes into a buffer so madevent can use the tables too The definitions could not be read from a table under helicity recycling: that rewriting indexes AMP by helicity as AMP(NCOMB,NGRAPHS), so the definitions cannot sit at the end of it and share one index with the amplitudes. Read the amplitudes of the current helicity into a buffer first, and run the definitions over that. The buffer is one dimensional whatever AMP looks like, so the table is the same one the standalone code uses, and the definitions stay one dimensional instead of gaining a helicity index they would never use. The recycler needs no change. It rewrites every AMP( it finds in the JAMP block to AMP( K,, and the only one left there is the gather itself, which is exactly what wants the helicity index: AMPBUF(ITMP) = AMP(ITMP) -> AMPBUF(ITMP) = AMP( K,ITMP) The declarations reach the rewritten file because the template it is built from comes out of the same replace_dict as the matrix element. The copy is bought back. On g g > 5g, with the JAMP block run once per helicity for 80 helicities, the definitions written out and read from AMP(K,i) strided take 1.5377 ms a call against 1.2312 for the gather, and gfortran -O2 takes 436 s against 0.17 s. The DATA statements have their own loop variable now, the templates already use I for something else. Checked with a full run of g g > g g g, tables forced on: every number of every results.dat is what main gives, only the timings differ. Co-Authored-By: Claude Opus 5 --- madgraph/iolibs/export_v4.py | 87 +++++++++++++------ .../matrix_madevent_group_v4.inc | 4 +- .../matrix_madevent_group_v4_hel.inc | 4 +- .../template_files/matrix_madevent_v4.inc | 4 +- .../matrix1.f | 4 +- .../matrix.f | 4 +- 6 files changed, 74 insertions(+), 33 deletions(-) diff --git a/madgraph/iolibs/export_v4.py b/madgraph/iolibs/export_v4.py index 12564564c8..0d9b815559 100755 --- a/madgraph/iolibs/export_v4.py +++ b/madgraph/iolibs/export_v4.py @@ -240,6 +240,11 @@ class ProcessExporterFortran(VirtualExporter): # finish with the plain scan once the orbit rounds have nothing left to # take as a whole (only used by the table emission, see below) jamp_greedy_tail = True + # Read the amplitudes of the current helicity into a buffer before running + # the definitions over it, instead of holding the definitions at the end of + # AMP. Needed where AMP is indexed by helicity, which is what madevent does + # once it rewrites the matrix element for helicity recycling. + jamp_gather = False # up to this many entries in the matrix, both optimisations are run and the # shorter result kept (see optimise_jamp_best) jamp_compare_max_size = 20000 @@ -2189,14 +2194,14 @@ def get_color_data_lines(self, matrix_element, n=128): return ret_list @staticmethod - def get_int_data_lines(name, values, n=128): + def get_int_data_lines(name, values, n=128, var='i'): """DATA statements filling the one dimensional integer array name.""" lines = [] for start in range(0, len(values), n): chunk = values[start:start + n] - lines.append(" DATA (%s(i),i=%d,%d) /%s/" % \ - (name, start + 1, start + len(chunk), + lines.append(" DATA (%s(%s),%s=%d,%d) /%s/" % \ + (name, var, var, start + 1, start + len(chunk), ','.join(str(int(v)) for v in chunk))) return lines @@ -2727,7 +2732,8 @@ def format(frac): self.jamp_recipes = recipes if recipes: - tmp_name = lambda k: "AMP(NGRAPHS+%d)" % k + buffer = self.jamp_buffer() + tmp_name = lambda k: "%s(NGRAPHS+%d)" % (buffer, k) defs = recipes['defs'] res_list.append("C The definitions below come in orbits of the") res_list.append("C permutations leaving the color basis") @@ -2742,6 +2748,14 @@ def format(frac): " + TMP_JAMP_F(ITMP)*AMP(TMP_JAMP_B(ITMP))") res_list.append(" ENDDO") else: + if self.jamp_gather: + res_list.append("C the amplitudes of this helicity are") + res_list.append("C read into one array first, so that") + res_list.append("C the definitions below take both") + res_list.append("C their operands from the same place") + res_list.append(" DO ITMP = 1, NGRAPHS") + res_list.append(" %s(ITMP) = AMP(ITMP)" % buffer) + res_list.append(" ENDDO") res_list.append("C the definitions of one level use none") res_list.append("C of each other, so they are sorted by") res_list.append("C the factor in front of the second") @@ -2750,20 +2764,23 @@ def format(frac): res_list.append(" DO ILEV = 1, NB_LEVEL") res_list.append(" DO ITMP = TMP_JAMP_L(5*ILEV-4)," " TMP_JAMP_L(5*ILEV-3)") - res_list.append(" AMP(NGRAPHS+ITMP) = AMP(TMP_JAMP_A(ITMP))" - " + AMP(TMP_JAMP_B(ITMP))") + res_list.append(" %s(NGRAPHS+ITMP) = %s(TMP_JAMP_A(ITMP))" + " + %s(TMP_JAMP_B(ITMP))" + % (buffer, buffer, buffer)) res_list.append(" ENDDO") res_list.append(" DO ITMP = TMP_JAMP_L(5*ILEV-3)+1," " TMP_JAMP_L(5*ILEV-2)") - res_list.append(" AMP(NGRAPHS+ITMP) = AMP(TMP_JAMP_A(ITMP))" - " - AMP(TMP_JAMP_B(ITMP))") + res_list.append(" %s(NGRAPHS+ITMP) = %s(TMP_JAMP_A(ITMP))" + " - %s(TMP_JAMP_B(ITMP))" + % (buffer, buffer, buffer)) res_list.append(" ENDDO") res_list.append(" DO ITMP = TMP_JAMP_L(5*ILEV-2)+1," " TMP_JAMP_L(5*ILEV-1)") - res_list.append(" AMP(NGRAPHS+ITMP) = AMP(TMP_JAMP_A(ITMP))" + res_list.append(" %s(NGRAPHS+ITMP) = %s(TMP_JAMP_A(ITMP))" " + TMP_JAMP_F(TMP_JAMP_L(5*ILEV)+ITMP" "-TMP_JAMP_L(5*ILEV-2))" - "*AMP(TMP_JAMP_B(ITMP))") + "*%s(TMP_JAMP_B(ITMP))" + % (buffer, buffer, buffer)) res_list.append(" ENDDO") res_list.append(" ENDDO") else: @@ -2789,7 +2806,8 @@ def format(frac): max_jamp=0 for (jamp, var), factor in new_mat.items(): if var > 0: - name = AMP_format % var + name = ("%s(%%s)" % self.jamp_buffer()) % var \ + if (recipes and self.jamp_gather) else AMP_format % var else: if recipes and recipes.get('factor_of'): # the definitions were renumbered, and one of them can be @@ -3457,14 +3475,14 @@ def jamp_definition_levels(defs): return [levels[key] for key in sorted(levels)] @staticmethod - def jamp_number_data_lines(name, values, per_line): + def jamp_number_data_lines(name, values, per_line, var='IJMP'): """DATA statements filling one array with the given constants.""" lines = [] for start in range(0, len(values), per_line): chunk = values[start:start + per_line] - lines.append(" DATA (%s(i),i=%d,%d) /%s/" % - (name, start + 1, start + len(chunk), + lines.append(" DATA (%s(%s),%s=%d,%d) /%s/" % + (name, var, var, start + 1, start + len(chunk), ','.join(chunk))) return lines @@ -3492,8 +3510,8 @@ def where(column): bounds = recipes['bounds'] lines = [ " INTEGER ITMP, ILEV", - "C I is the loop variable of the DATA statements below", - " INTEGER I", + "C IJMP is the loop variable of the DATA statements below", + " INTEGER IJMP", " INTEGER NB_TMP_JAMP, NB_LEVEL, NB_GENERAL", " PARAMETER (NB_TMP_JAMP=%d)" % nb_def, " PARAMETER (NB_LEVEL=%d)" % len(bounds), @@ -3502,10 +3520,15 @@ def where(column): " INTEGER TMP_JAMP_L(5*NB_LEVEL)", " %s TMP_JAMP_F(NB_GENERAL)" % self.jamp_factor_type(recipes), ] - lines += self.get_int_data_lines("TMP_JAMP_A", left) - lines += self.get_int_data_lines("TMP_JAMP_B", right) + if self.jamp_gather: + lines.append(" COMPLEX*16 AMPBUF(%d+NB_TMP_JAMP)" % nb_amp) + lines += self.get_int_data_lines("TMP_JAMP_A", left, + var="IJMP") + lines += self.get_int_data_lines("TMP_JAMP_B", right, + var="IJMP") lines += self.get_int_data_lines("TMP_JAMP_L", - sum((list(b) for b in bounds), [])) + sum((list(b) for b in bounds), []), + var="IJMP") assert len(bounds[0]) == 5 if general: lines += self.jamp_number_data_lines("TMP_JAMP_F", factor, @@ -3524,6 +3547,11 @@ def jamp_number(value): return "%dD0" % int(value) return ("%.15e" % value).replace('e', 'd') + def jamp_buffer(self): + """Array the definitions and their operands are read from.""" + + return 'AMPBUF' if self.jamp_gather else 'AMP' + @staticmethod def jamp_power_data(recipes): """DATA statement for the four powers of i, real when none of them is @@ -3729,16 +3757,11 @@ def get_jamp_init_routine(self, recipes, proc_prefix): def jamp_tables_allowed(self): """Whether the definitions may be read from a table rather than written - out. That needs the template to declare the arrays and to hold the - definitions at the end of AMP, which only the standalone one does. - - Madevent cannot: with helicity recycling on, which is its default, the - matrix element is rewritten from a template of its own where AMP is - indexed by helicity as AMP(NCOMB,NGRAPHS) and only the JAMP lines are - carried over. The definitions are written out there, and only the - optimisation itself is shared.""" + out. Both the standalone and the madevent templates declare what that + needs; madevent reads the amplitudes of the current helicity into a + buffer first, see jamp_gather.""" - return not isinstance(self, ProcessExporterFortranME) + return True def jamp_orbit_allowed(self, matrix_element): """Whether the orbit equivariant optimisation is used here.""" @@ -6780,6 +6803,9 @@ class ProcessExporterFortranME(ProcessExporterFortran): matrix_file = "matrix_madevent_v4.inc" jamp_orbit = True + # AMP is indexed by helicity once the matrix element is rewritten for + # helicity recycling, so the definitions cannot sit at the end of it + jamp_gather = True done_warning_tchannel = False default_opt = {'clean': False, 'complex_mass':False, @@ -7566,6 +7592,11 @@ def write_matrix_element_v4(self, writer, matrix_element, fortran_model, orbit=self.jamp_orbit_allowed(matrix_element)) replace_dict['jamp_lines'] = '\n'.join(jamp_lines) replace_dict['nb_temp_jamp'] = nb_temp + recipes = getattr(self, 'jamp_recipes', None) + replace_dict['jamp_decl'] = '\n'.join( + self.get_jamp_decl_lines(recipes, '')) + replace_dict['jamp_tmp_decl'] = '' if recipes else \ + " COMPLEX*16 TMP_JAMP(%i)" % nb_temp if self.beam_polarization == [True, True]: replace_dict['beam_polarization'] = """ diff --git a/madgraph/iolibs/template_files/matrix_madevent_group_v4.inc b/madgraph/iolibs/template_files/matrix_madevent_group_v4.inc index b1523fc0e2..eb252d191d 100644 --- a/madgraph/iolibs/template_files/matrix_madevent_group_v4.inc +++ b/madgraph/iolibs/template_files/matrix_madevent_group_v4.inc @@ -321,7 +321,9 @@ C C LOCAL VARIABLES C INTEGER I,J,M,N - COMPLEX*16 ZTEMP, TMP_JAMP(%(nb_temp_jamp)i) + COMPLEX*16 ZTEMP +%(jamp_tmp_decl)s +%(jamp_decl)s INTEGER CF(NCOLOR*(NCOLOR+1)/2) INTEGER DENOM, CF_INDEX COMPLEX*16 AMP(NGRAPHS), JAMP(NCOLOR,NAMPSO) diff --git a/madgraph/iolibs/template_files/matrix_madevent_group_v4_hel.inc b/madgraph/iolibs/template_files/matrix_madevent_group_v4_hel.inc index 45bcdb8211..5cdb43bbe1 100644 --- a/madgraph/iolibs/template_files/matrix_madevent_group_v4_hel.inc +++ b/madgraph/iolibs/template_files/matrix_madevent_group_v4_hel.inc @@ -234,7 +234,9 @@ C C LOCAL VARIABLES C INTEGER I,J,M,N,K - COMPLEX*16 ZTEMP,TMP_JAMP(%(nb_temp_jamp)i) + COMPLEX*16 ZTEMP +%(jamp_tmp_decl)s +%(jamp_decl)s COMPLEX*16 TMP(%(wavefunctionsize)d) INTEGER CF(NCOLOR*(NCOLOR+1)) INTEGER DENOM, CF_INDEX diff --git a/madgraph/iolibs/template_files/matrix_madevent_v4.inc b/madgraph/iolibs/template_files/matrix_madevent_v4.inc index 7b1d6bf1b9..a78cfbe917 100644 --- a/madgraph/iolibs/template_files/matrix_madevent_v4.inc +++ b/madgraph/iolibs/template_files/matrix_madevent_v4.inc @@ -271,7 +271,9 @@ C C LOCAL VARIABLES C INTEGER I,J,M,N - COMPLEX*16 ZTEMP, TMP_JAMP(%(nb_temp_jamp)i) + COMPLEX*16 ZTEMP +%(jamp_tmp_decl)s +%(jamp_decl)s INTEGER CF(NCOLOR*(NCOLOR+1)) INTEGER CF_INDEX,DENOM COMPLEX*16 AMP(NGRAPHS), JAMP(NCOLOR,NAMPSO) diff --git a/tests/input_files/IOTestsComparison/IOExportV4IOTest/export_matrix_element_v4_madevent_group/matrix1.f b/tests/input_files/IOTestsComparison/IOExportV4IOTest/export_matrix_element_v4_madevent_group/matrix1.f index 9805b01974..25bc11f65a 100644 --- a/tests/input_files/IOTestsComparison/IOExportV4IOTest/export_matrix_element_v4_madevent_group/matrix1.f +++ b/tests/input_files/IOTestsComparison/IOExportV4IOTest/export_matrix_element_v4_madevent_group/matrix1.f @@ -397,7 +397,9 @@ REAL*8 FUNCTION MATRIX1(P,NHEL,IFLAV, IHEL,AMP2, JAMP2, IVEC) C LOCAL VARIABLES C INTEGER I,J,M,N - COMPLEX*16 ZTEMP, TMP_JAMP(0) + COMPLEX*16 ZTEMP + COMPLEX*16 TMP_JAMP(0) + INTEGER CF(NCOLOR*(NCOLOR+1)/2) INTEGER DENOM, CF_INDEX COMPLEX*16 AMP(NGRAPHS), JAMP(NCOLOR,NAMPSO) diff --git a/tests/input_files/IOTestsComparison/IOExportV4IOTest/export_matrix_element_v4_madevent_nogroup/matrix.f b/tests/input_files/IOTestsComparison/IOExportV4IOTest/export_matrix_element_v4_madevent_nogroup/matrix.f index c92934f27b..3ee6679790 100644 --- a/tests/input_files/IOTestsComparison/IOExportV4IOTest/export_matrix_element_v4_madevent_nogroup/matrix.f +++ b/tests/input_files/IOTestsComparison/IOExportV4IOTest/export_matrix_element_v4_madevent_nogroup/matrix.f @@ -332,7 +332,9 @@ REAL*8 FUNCTION MATRIX(P,NHEL,IFLAV, IVEC) C LOCAL VARIABLES C INTEGER I,J,M,N - COMPLEX*16 ZTEMP, TMP_JAMP(0) + COMPLEX*16 ZTEMP + COMPLEX*16 TMP_JAMP(0) + INTEGER CF(NCOLOR*(NCOLOR+1)) INTEGER CF_INDEX,DENOM COMPLEX*16 AMP(NGRAPHS), JAMP(NCOLOR,NAMPSO) From 018e56b67800e51bb763239ce771249f8a023daf Mon Sep 17 00:00:00 2001 From: Olivier Mattelaer Date: Thu, 6 Aug 2026 21:58:34 +0200 Subject: [PATCH 12/42] find the reversal symmetry of the color basis and fold the color matrix on it Reversing every color basis element maps the basis onto itself, and for a pure gluon process the color coefficients of a line and of its reverse differ by one overall sign: JAMP[reverse(i)] = (-1)^n JAMP[i]. Half the color flows therefore carry nothing of their own, and |M|^2 can be summed over one line per pair with a folded color matrix instead of over every line. That is half the color flows and a quarter of the color sum, which is 78% of the run time of g g > 6g. This is the reading half of it, not yet used by the generated code: - reverse_immutable, the reversal of a basis key - get_jamp_reflection, which reads the sign off the color coefficients rather than assuming it, so a process where the relation does not hold gets None. It holds for pure gluons (+1 for an even number, -1 for an odd one) and not for a quark line, where reversing does not commute with the fermion flow. - jamp_folded_color_matrix, C'[a][b] summed over the two lines of each pair with their signs. Checked against the unfolded contraction for g g > g g, g g > 3g and g g > 4g: same |M|^2 to all digits, same denominator, exactly half the lines and no self-paired line in any of them. Co-Authored-By: Claude Opus 5 --- madgraph/core/color_amp.py | 19 +++++++ madgraph/iolibs/export_v4.py | 96 ++++++++++++++++++++++++++++++++++++ 2 files changed, 115 insertions(+) diff --git a/madgraph/core/color_amp.py b/madgraph/core/color_amp.py index 9250f1a69f..04a27c6ee2 100755 --- a/madgraph/core/color_amp.py +++ b/madgraph/core/color_amp.py @@ -632,6 +632,25 @@ def permute_immutable(struct, perm): return tuple(res) +def reverse_immutable(struct): + """Reverse every color object of an immutable color basis key, bringing the + result back to the canonical form. A trace is cyclic, so it is rotated onto + its smallest index afterwards. Returns None for anything which is not built + of traces alone, which is where the reversal has a meaning of its own.""" + + res = [] + for name, indices in struct: + if name != 'Tr': + return None + indices = tuple(reversed(indices)) + if len(indices) > 1: + start = min(range(len(indices)), key=indices.__getitem__) + indices = indices[start:] + indices[:start] + res.append((name, indices)) + res.sort() + return tuple(res) + + class ColorBasisSymmetry(object): """Permutations of the external color indices which map a color basis (or a pair of color bases, for an asymmetric color matrix) onto itself. diff --git a/madgraph/iolibs/export_v4.py b/madgraph/iolibs/export_v4.py index 0d9b815559..47481f6dcf 100755 --- a/madgraph/iolibs/export_v4.py +++ b/madgraph/iolibs/export_v4.py @@ -3034,6 +3034,102 @@ def optimise_jamp_best(self, all_element, symmetry): # does, leaves the matrix invariant at every step, and the definitions can # be written as one recipe per orbit. + def get_jamp_reflection(self, matrix_element, all_element): + """Reversing every color basis element maps the basis onto itself, and + for a pure gluon process the color coefficients of a line and of its + reverse differ by one overall sign, so half the color flows carry no + information of their own: + + JAMP[reverse(i)] = sign * JAMP[i] + + Return (reverse, sign) or None. The relation is read off the color + coefficients themselves rather than assumed, so a process where it does + not hold -- a quark line, where reversing does not commute with the + fermion flow -- simply gets None.""" + + if not isinstance(matrix_element, helas_objects.HelasMatrixElement): + return None + color_basis = matrix_element.get('color_basis') + if not color_basis or len(color_basis) < 2: + return None + keys = sorted(color_basis.keys()) + position = dict((key, i) for i, key in enumerate(keys)) + + reverse = [] + for key in keys: + other = color_amp.reverse_immutable(key) + if other is None or other not in position: + return None + reverse.append(position[other]) + if any(reverse[reverse[i]] != i for i in range(len(keys))): + return None + + columns = collections.defaultdict(dict) + for (i, j), value in all_element.items(): + if value: + columns[i][j] = value + + sign = None + for i in range(len(keys)): + here, there = columns.get(i + 1, {}), columns.get(reverse[i] + 1, {}) + if set(here) != set(there): + return None + for amp, value in here.items(): + ratio = there[amp] / value + if ratio not in (1, -1): + return None + if sign is None: + sign = int(ratio.real) + elif sign != int(ratio.real): + return None + if sign is None: + return None + return reverse, sign + + def jamp_folded_color_matrix(self, matrix_element, reverse, sign): + """The color matrix over one line per reversal pair. Summing |M|^2 over + the pairs instead of over every line gives the same number, since the + two lines of a pair only differ by the overall sign: + + C'[a][b] = sum over the two lines of a and the two of b, each + weighted by its sign relative to the line kept + + Returns (denominator, rows) with rows[a][b] integer, a and b indexing + the representatives.""" + + color_matrix = matrix_element.get('color_matrix') + representatives, _slot = self.jamp_reflection_representatives(reverse) + denominator = max(color_matrix.get_line_denominators()) + full = [color_matrix.get_line_numerators(i, denominator) + for i in range(len(reverse))] + + def pair(a): + return [(a, 1)] if reverse[a] == a else [(a, 1), (reverse[a], sign)] + + rows = [] + for a in representatives: + row = [] + for b in representatives: + total = 0 + for i, ci in pair(a): + for j, cj in pair(b): + total += ci * cj * full[i][j] + assert int(total) == total + row.append(int(total)) + rows.append(row) + return denominator, rows + + @staticmethod + def jamp_reflection_representatives(reverse): + """One line per pair, and for every line the pair it belongs to.""" + + representatives = [i for i in range(len(reverse)) if i <= reverse[i]] + slot = {} + for index, line in enumerate(representatives): + slot[line] = index + slot[reverse[line]] = index + return representatives, slot + @staticmethod def jamp_column_form(column): """Canonical form of one column of the JAMP matrix up to a global sign, From 04ce4e06c448950f2b4db747688cfbc723570db4 Mon Sep 17 00:00:00 2001 From: Olivier Mattelaer Date: Thu, 6 Aug 2026 22:37:14 +0200 Subject: [PATCH 13/42] sum |M|^2 over one color flow per reversal pair Reversing a color flow gives another flow of the basis whose amplitude is the same up to one overall sign, so half of them carry nothing of their own. Sum over one per pair against a color matrix folded onto those, and the color sum does a quarter of the work. It is 78% of the run time of g g > 6g, so this is the largest single lever there is short of changing the color basis itself. GET_MATRIX gathers the flows it keeps out of JAMP through COLREP and runs the same loop over them. JAMP stays the full length, so the color flow decomposition, GET_INTER and the density matrix see exactly what they saw before. Two ways of getting the folded matrix to the generated code: - sign +1: reversing commutes with permuting the color indices, so a permutation of the flows is also a permutation of the pairs and the folded matrix still has one line per orbit. It is rebuilt at run time as before. - sign -1: a permutation may send a flow onto its own partner, which flips its weight, and the rebuilt form has nowhere to put that sign. The folded matrix is written out instead, which is affordable while it stays small (g g > 5g folds to 360 lines, 65k entries) and folding is declined above color_fold_max_written. The sign is read off the color coefficients rather than assumed, so a quark line, where reversing does not commute with the fermion flow, simply does not fold: g g > t t~ g g and u u~ > u u~ g g are untouched and their |M|^2 is unchanged to the last digit. g g > 5g: GET_MATRIX 0.3459 -> 0.0844 ms a call, 4.10x, the 4x expected. |M|^2 6.6739867626784571E-007 -> ...550E-007, 15 digits: the folded sum regroups 720^2 terms into 360^2 with larger coefficients, so it rounds differently. g g > 4g and g g > 3g are unchanged to the last digit. test_generate_helas_diagrams_gg_gg checks the emitted color matrix against a hand written 6x6, which is the unfolded one; it now asks for that explicitly. Co-Authored-By: Claude Opus 5 --- madgraph/iolibs/export_v4.py | 199 ++++++++++++++++-- .../template_files/matrix_standalone_v4.inc | 28 ++- .../matrix.f | 31 ++- tests/unit_tests/iolibs/test_export_v4.py | 4 + 4 files changed, 224 insertions(+), 38 deletions(-) diff --git a/madgraph/iolibs/export_v4.py b/madgraph/iolibs/export_v4.py index 47481f6dcf..c1c1fb2a83 100755 --- a/madgraph/iolibs/export_v4.py +++ b/madgraph/iolibs/export_v4.py @@ -229,6 +229,8 @@ class ProcessExporterFortran(VirtualExporter): jamp_optim = False # how many times the JAMP optimisation called itself, for the record myjamp_count = 0 + # sum |M|^2 over one color flow per reversal pair instead of over every one + jamp_fold = True # write the JAMP definitions as one recipe per orbit of the permutations # leaving the color basis invariant, instead of one line per definition jamp_orbit = False @@ -2094,6 +2096,49 @@ def format_integer_list(self, list, name, n=5): + @staticmethod + def jamp_fold_spanning_tree(permutations, size): + """Same walk as ColorBasisSymmetry.spanning_tree, over an index set + given directly as permutations rather than as color basis keys.""" + + keep = [] + parent_uf = list(range(size)) + + def find(x): + while parent_uf[x] != x: + parent_uf[x] = parent_uf[parent_uf[x]] + x = parent_uf[x] + return x + + for perm in permutations: + used = False + for i, j in enumerate(perm): + ri, rj = find(i), find(j) + if ri != rj: + parent_uf[ri] = rj + used = True + if used: + keep.append(perm) + + representative = [-1] * size + parent = [None] * size + representatives = [] + for start in range(size): + if representative[start] != -1: + continue + representatives.append(start) + representative[start] = start + queue = collections.deque([start]) + while queue: + current = queue.popleft() + for local, perm in enumerate(keep): + image = perm[current] + if representative[image] == -1: + representative[image] = start + parent[image] = (current, local) + queue.append(image) + return representatives, representative, parent, keep + def get_color_matrix_encoding(self, matrix_element): """Describe the color matrix by one line per orbit of the index permutations leaving the color basis invariant, plus the permutations @@ -2116,7 +2161,27 @@ def get_color_matrix_encoding(self, matrix_element): symmetry = color_amp.ColorBasisSymmetry(keys) if not symmetry.has_symmetry(): return None - representatives, representative, parent, gens = symmetry.spanning_tree() + + folding = self.get_jamp_folding(matrix_element) + if folding and folding['sign'] < 0: + # The rebuilt form cannot carry the weight a permutation picks up + # when it sends a line onto its own partner. The sum runs over the + # folded matrix either way, so there is no falling back to the + # unfolded encoding here: it has to be written out instead. + return None + if folding: + # reversing commutes with permuting the indices, so a permutation + # of the lines is also a permutation of the pairs + slot = folding['slot'] + nb_color = len(folding['representatives']) + induced = [[slot[perm[line]] + for line in folding['representatives']] + for perm in symmetry.generators1] + representatives, representative, parent, gens = \ + self.jamp_fold_spanning_tree(induced, nb_color) + else: + representatives, representative, parent, gens = \ + symmetry.spanning_tree() # Writing the entries out is well trodden and the compressed form # carries a routine of its own, so only take it when it pays clearly. @@ -2127,13 +2192,18 @@ def get_color_matrix_encoding(self, matrix_element): if size * self.color_encoding_margin > nb_color * (nb_color + 1) // 2: return None - denominator = max(color_matrix.get_line_denominators()) - slot = dict((line, index) for index, line in enumerate(representatives)) - rows = [] - for line in representatives: - num_list = color_matrix.get_line_numerators(line, denominator) - assert all(int(i) == i for i in num_list) - rows.append([int(i) for i in num_list]) + place = dict((line, index) for index, line in enumerate(representatives)) + if folding: + denominator, folded = self.jamp_folded_color_matrix( + matrix_element, folding['reverse'], folding['sign']) + rows = [folded[line] for line in representatives] + else: + denominator = max(color_matrix.get_line_denominators()) + rows = [] + for line in representatives: + num_list = color_matrix.get_line_numerators(line, denominator) + assert all(int(i) == i for i in num_list) + rows.append([int(i) for i in num_list]) return {'denom': denominator, 'nb_color': nb_color, @@ -2143,7 +2213,7 @@ def get_color_matrix_encoding(self, matrix_element): # reaching it, or (0,0) when the line is a representative 'parent': [(0, 0) if p is None else (p[0] + 1, p[1] + 1) for p in parent], - 'slot': [slot[representative[i]] + 1 + 'slot': [place[representative[i]] + 1 for i in range(nb_color)]} def get_color_data_lines(self, matrix_element, n=128): @@ -2161,6 +2231,26 @@ def get_color_data_lines(self, matrix_element, n=128): return ["DATA %%(proc_prefix)sDenom/%(denom)i/" % \ {'denom': denominator}] + folding = self.get_jamp_folding(matrix_element) + if folding: + denominator, folded = self.jamp_folded_color_matrix( + matrix_element, folding['reverse'], folding['sign']) + ret_list = ["DATA %%(proc_prefix)sDenom/%(denom)i/" % + {'denom': denominator}] + cf_index = 0 + for index in range(len(folded)): + row = folded[index] + for k in range(index, len(row), n): + chunk = row[k:k + n] + ret_list.append( + "DATA (%%(proc_prefix)sCF(i),i=%3r,%3r) /%s/" % + (cf_index + 1, cf_index + len(chunk), + ','.join("%i" % ((1 if (k == index and pos == 0) + else 2) * int(v)) + for pos, v in enumerate(chunk)))) + cf_index += len(chunk) + return ret_list + ret_list = [] my_cs = color.ColorString() denominator = max(matrix_element.get('color_matrix').get_line_denominators()) @@ -2210,8 +2300,9 @@ def get_color_init_routine(self, matrix_element, proc_prefix): description, or an empty routine when the entries are written out.""" encoding = self.get_color_matrix_encoding(matrix_element) - nb_color = len(matrix_element.get('color_matrix')._sorted_keys1) \ - if matrix_element.get('color_matrix') else 0 + nb_color = encoding['nb_color'] if encoding else \ + (len(matrix_element.get('color_matrix')._sorted_keys1) + if matrix_element.get('color_matrix') else 0) header = [" SUBROUTINE %sINIT_CF()" % proc_prefix] if not encoding: return header + [" RETURN", " END"] @@ -3034,7 +3125,30 @@ def optimise_jamp_best(self, all_element, symmetry): # does, leaves the matrix invariant at every step, and the definitions can # be written as one recipe per orbit. - def get_jamp_reflection(self, matrix_element, all_element): + @staticmethod + def jamp_color_rows(matrix_element): + """The color coefficient of every amplitude, one dictionary per color + basis line. Same numbers get_JAMP_lines works from.""" + + rows = [] + powers = {} + for coeff_list in matrix_element.get_color_amplitudes(): + row = {} + for coefficient, amp in coeff_list: + if not coefficient: + continue + try: + power = powers[coefficient[3]] + except KeyError: + power = fractions.Fraction(3) ** coefficient[3] + powers[coefficient[3]] = power + value = (1j if coefficient[2] else 1) * coefficient[0] * \ + coefficient[1] * power + row[amp] = row.get(amp, 0) + value + rows.append(dict((amp, complex(v)) for amp, v in row.items() if v)) + return rows + + def get_jamp_reflection(self, matrix_element): """Reversing every color basis element maps the basis onto itself, and for a pure gluon process the color coefficients of a line and of its reverse differ by one overall sign, so half the color flows carry no @@ -3064,14 +3178,11 @@ def get_jamp_reflection(self, matrix_element, all_element): if any(reverse[reverse[i]] != i for i in range(len(keys))): return None - columns = collections.defaultdict(dict) - for (i, j), value in all_element.items(): - if value: - columns[i][j] = value + columns = self.jamp_color_rows(matrix_element) sign = None for i in range(len(keys)): - here, there = columns.get(i + 1, {}), columns.get(reverse[i] + 1, {}) + here, there = columns[i], columns[reverse[i]] if set(here) != set(there): return None for amp, value in here.items(): @@ -3086,6 +3197,36 @@ def get_jamp_reflection(self, matrix_element, all_element): return None return reverse, sign + # Above this many entries the folded color matrix is not written out but + # rebuilt at run time, which only the sign +1 case can do (see + # get_jamp_folding). + color_fold_max_written = 300000 + + def get_jamp_folding(self, matrix_element): + """Whether to sum |M|^2 over one line per reversal pair, and the + (reverse, sign, representatives, slot) that goes with it. + + With sign +1 every line of a pair enters with the same weight, so the + permutations leaving the color basis invariant carry over to the pairs + unchanged and the folded matrix can still be rebuilt at run time from + one line per orbit. With sign -1 a permutation may send a line onto its + own partner, which flips the weight, and the rebuilt form would need a + sign of its own; there the folded matrix is written out instead, which + is only affordable while it stays small.""" + + if not self.jamp_fold: + return None + found = self.get_jamp_reflection(matrix_element) + if not found: + return None + reverse, sign = found + representatives, slot = self.jamp_reflection_representatives(reverse) + nb = len(representatives) + if sign < 0 and nb * (nb + 1) // 2 > self.color_fold_max_written: + return None + return {'reverse': reverse, 'sign': sign, + 'representatives': representatives, 'slot': slot} + def jamp_folded_color_matrix(self, matrix_element, reverse, sign): """The color matrix over one line per reversal pair. Summing |M|^2 over the pairs instead of over every line gives the same number, since the @@ -5646,7 +5787,29 @@ def write_matrix_element_v4(self, writer, matrix_element, fortran_model, # Extract ncolor ncolor = max(1, len(matrix_element.get('color_basis'))) replace_dict['ncolor'] = ncolor - replace_dict['ncolortriang'] = ncolor * (ncolor + 1) // 2 + # |M|^2 is summed over one color flow per reversal pair when the basis + # allows it, so the color matrix is only over those + folding = self.get_jamp_folding(matrix_element) + self.jamp_folding = folding + nfold = len(folding['representatives']) if folding else ncolor + replace_dict['ncolorfold'] = nfold + replace_dict['ncolortriang'] = nfold * (nfold + 1) // 2 + replace_dict['color_fold_index'] = '\n'.join( + self.get_int_data_lines("COLREP", + [i + 1 for i in folding['representatives']], + var='ICF')) if folding else '' + replace_dict['color_fold_gather'] = ( + " DO ICF = 1, NCOLORFOLD\n" + " JFOLD(ICF) = JAMP(COLREP(ICF))\n" + " ENDDO" if folding else + " DO ICF = 1, NCOLOR\n" + " JFOLD(ICF) = JAMP(ICF)\n" + " ENDDO") + if not folding: + replace_dict['color_fold_decl'] = " INTEGER ICF" + if folding: + replace_dict['color_fold_decl'] = \ + " INTEGER COLREP(NCOLORFOLD)\n INTEGER ICF" replace_dict['hel_avg_factor'] = matrix_element.get_hel_avg_factor() replace_dict['beamone_helavgfactor'], replace_dict['beamtwo_helavgfactor'] =\ diff --git a/madgraph/iolibs/template_files/matrix_standalone_v4.inc b/madgraph/iolibs/template_files/matrix_standalone_v4.inc index d8239ac9e3..db5f614509 100644 --- a/madgraph/iolibs/template_files/matrix_standalone_v4.inc +++ b/madgraph/iolibs/template_files/matrix_standalone_v4.inc @@ -364,8 +364,9 @@ CF2PY INTENT(OUT) :: MATRIX CF2PY INTENT(IN) :: JAMP - INTEGER NCOLOR + INTEGER NCOLOR, NCOLORFOLD PARAMETER (NCOLOR=%(ncolor)d) + PARAMETER (NCOLORFOLD=%(ncolorfold)d) REAL*8 ZERO,MATRIX PARAMETER (ZERO=0D0) C @@ -376,10 +377,13 @@ C COMPLEX*16 ZTEMP,Z1,Z2,Z3,Z4 INTEGER CF_INDEX - INTEGER %(proc_prefix)sCF(NCOLOR*(NCOLOR+1)/2) + INTEGER %(proc_prefix)sCF(NCOLORFOLD*(NCOLORFOLD+1)/2) INTEGER %(proc_prefix)sDENOM common /%(proc_prefix)scolor_matrix/ %(proc_prefix)sCF,%(proc_prefix)sDENOM COMPLEX*16 JAMP(NCOLOR) + COMPLEX*16 JFOLD(NCOLORFOLD) +%(color_fold_decl)s +%(color_fold_index)s %(jamp_tmp_decl)s COMPLEX*16 DUM0,DUM1 DATA DUM0, DUM1/(0d0, 0d0), (1d0, 0d0)/ @@ -389,6 +393,10 @@ C COLOR DATA C CALL %(proc_prefix)sINIT_CF() +C Reversing a color flow gives the same one back up to an overall +C sign, so only one of each pair carries anything: the sum below runs +C over those, against a color matrix folded onto them. +%(color_fold_gather)s MATRIX = 0.D0 CF_INDEX = 0 C Four accumulators, not one: with a single one every @@ -396,25 +404,25 @@ C term waits for the one before it to come out of the C adder, and that latency is what the loop spends its C time on. No compiler does this by itself, since it C changes the order the terms are summed in. - DO I = 1, NCOLOR + DO I = 1, NCOLORFOLD Z1 = (0.D0,0.D0) Z2 = (0.D0,0.D0) Z3 = (0.D0,0.D0) Z4 = (0.D0,0.D0) - NJ = NCOLOR - I + 1 + NJ = NCOLORFOLD - I + 1 NB = (NJ/4)*4 DO J = 0, NB-4, 4 - Z1 = Z1 + %(proc_prefix)sCF(CF_INDEX+J+1)*JAMP(I+J) - Z2 = Z2 + %(proc_prefix)sCF(CF_INDEX+J+2)*JAMP(I+J+1) - Z3 = Z3 + %(proc_prefix)sCF(CF_INDEX+J+3)*JAMP(I+J+2) - Z4 = Z4 + %(proc_prefix)sCF(CF_INDEX+J+4)*JAMP(I+J+3) + Z1 = Z1 + %(proc_prefix)sCF(CF_INDEX+J+1)*JFOLD(I+J) + Z2 = Z2 + %(proc_prefix)sCF(CF_INDEX+J+2)*JFOLD(I+J+1) + Z3 = Z3 + %(proc_prefix)sCF(CF_INDEX+J+3)*JFOLD(I+J+2) + Z4 = Z4 + %(proc_prefix)sCF(CF_INDEX+J+4)*JFOLD(I+J+3) ENDDO ZTEMP = (Z1+Z2)+(Z3+Z4) DO J = NB, NJ-1 - ZTEMP = ZTEMP + %(proc_prefix)sCF(CF_INDEX+J+1)*JAMP(I+J) + ZTEMP = ZTEMP + %(proc_prefix)sCF(CF_INDEX+J+1)*JFOLD(I+J) ENDDO CF_INDEX = CF_INDEX + NJ - MATRIX = MATRIX+ZTEMP*DCONJG(JAMP(I))/%(proc_prefix)sDENOM + MATRIX = MATRIX+ZTEMP*DCONJG(JFOLD(I))/%(proc_prefix)sDENOM ENDDO END diff --git a/tests/input_files/IOTestsComparison/IOExportV4IOTest/export_matrix_element_v4_standalone/matrix.f b/tests/input_files/IOTestsComparison/IOExportV4IOTest/export_matrix_element_v4_standalone/matrix.f index b8db276219..2a2368e94a 100644 --- a/tests/input_files/IOTestsComparison/IOExportV4IOTest/export_matrix_element_v4_standalone/matrix.f +++ b/tests/input_files/IOTestsComparison/IOExportV4IOTest/export_matrix_element_v4_standalone/matrix.f @@ -483,8 +483,9 @@ SUBROUTINE GET_MATRIX(JAMP,MATRIX) CF2PY INTENT(IN) :: JAMP - INTEGER NCOLOR + INTEGER NCOLOR, NCOLORFOLD PARAMETER (NCOLOR=1) + PARAMETER (NCOLORFOLD=1) REAL*8 ZERO,MATRIX PARAMETER (ZERO=0D0) C @@ -495,10 +496,13 @@ SUBROUTINE GET_MATRIX(JAMP,MATRIX) COMPLEX*16 ZTEMP,Z1,Z2,Z3,Z4 INTEGER CF_INDEX - INTEGER CF(NCOLOR*(NCOLOR+1)/2) + INTEGER CF(NCOLORFOLD*(NCOLORFOLD+1)/2) INTEGER DENOM COMMON /COLOR_MATRIX/ CF,DENOM COMPLEX*16 JAMP(NCOLOR) + COMPLEX*16 JFOLD(NCOLORFOLD) + INTEGER ICF + COMPLEX*16 TMP_JAMP(0) COMPLEX*16 DUM0,DUM1 DATA DUM0, DUM1/(0D0, 0D0), (1D0, 0D0)/ @@ -508,6 +512,13 @@ SUBROUTINE GET_MATRIX(JAMP,MATRIX) C CALL INIT_CF() +C Reversing a color flow gives the same one back up to an overall +C sign, so only one of each pair carries anything: the sum below +C runs +C over those, against a color matrix folded onto them. + DO ICF = 1, NCOLOR + JFOLD(ICF) = JAMP(ICF) + ENDDO MATRIX = 0.D0 CF_INDEX = 0 C Four accumulators, not one: with a single one every @@ -515,25 +526,25 @@ SUBROUTINE GET_MATRIX(JAMP,MATRIX) C adder, and that latency is what the loop spends its C time on. No compiler does this by itself, since it C changes the order the terms are summed in. - DO I = 1, NCOLOR + DO I = 1, NCOLORFOLD Z1 = (0.D0,0.D0) Z2 = (0.D0,0.D0) Z3 = (0.D0,0.D0) Z4 = (0.D0,0.D0) - NJ = NCOLOR - I + 1 + NJ = NCOLORFOLD - I + 1 NB = (NJ/4)*4 DO J = 0, NB-4, 4 - Z1 = Z1 + CF(CF_INDEX+J+1)*JAMP(I+J) - Z2 = Z2 + CF(CF_INDEX+J+2)*JAMP(I+J+1) - Z3 = Z3 + CF(CF_INDEX+J+3)*JAMP(I+J+2) - Z4 = Z4 + CF(CF_INDEX+J+4)*JAMP(I+J+3) + Z1 = Z1 + CF(CF_INDEX+J+1)*JFOLD(I+J) + Z2 = Z2 + CF(CF_INDEX+J+2)*JFOLD(I+J+1) + Z3 = Z3 + CF(CF_INDEX+J+3)*JFOLD(I+J+2) + Z4 = Z4 + CF(CF_INDEX+J+4)*JFOLD(I+J+3) ENDDO ZTEMP = (Z1+Z2)+(Z3+Z4) DO J = NB, NJ-1 - ZTEMP = ZTEMP + CF(CF_INDEX+J+1)*JAMP(I+J) + ZTEMP = ZTEMP + CF(CF_INDEX+J+1)*JFOLD(I+J) ENDDO CF_INDEX = CF_INDEX + NJ - MATRIX = MATRIX+ZTEMP*DCONJG(JAMP(I))/DENOM + MATRIX = MATRIX+ZTEMP*DCONJG(JFOLD(I))/DENOM ENDDO END diff --git a/tests/unit_tests/iolibs/test_export_v4.py b/tests/unit_tests/iolibs/test_export_v4.py index 847ff6932b..34e2db413c 100644 --- a/tests/unit_tests/iolibs/test_export_v4.py +++ b/tests/unit_tests/iolibs/test_export_v4.py @@ -3798,6 +3798,10 @@ def test_generate_helas_diagrams_gg_gg(self): denom = 6 i = 0 + # the numbers above are the color matrix over every color flow; the + # folded form sums each reversal pair into one line and is checked + # against |M|^2 itself elsewhere + exporter.jamp_fold = False for data in exporter.get_color_data_lines(\ matrix_element): From 48ab222d8846e674a8ae9156229bb7d310fb6fc1 Mon Sep 17 00:00:00 2001 From: Olivier Mattelaer Date: Thu, 6 Aug 2026 23:08:25 +0200 Subject: [PATCH 14/42] work out one color flow per reversal pair and copy the other The optimisation is given only the flows the color sum keeps, and the rest are one assignment each. Sharing means the definitions do not halve with the lines -- they are reused across flows -- but there are fewer of them: g g > 4g goes from 795 to 645, 19% off, with |M|^2 unchanged to the last digit. Only with sign +1. There both lines of a pair weigh the same, so a permutation of the lines is a permutation of the pairs with nothing attached and what is left of the matrix is still symmetric enough for the orbit version to work on it. With sign -1 a permutation may send a kept line onto its partner, which flips the weight for that pair alone; the matrix is then not symmetric in the form optimise_jamp wants, the symmetry is lost and with it the orbit optimisation and the table emission, which costs far more than the halving saves -- g g > 5g went to 41175 lines from 21575 when I tried. get_jamp_halving is where that is decided, and g g > 3g and g g > 5g keep every flow. The copies are still written: only the color sum skips those flows, while the color flow decomposition, GET_INTER and the density matrix read the whole array. g g > 3g, g g > t t~ g g and u u~ > u u~ g g are untouched, all 122 tests pass. Co-Authored-By: Claude Opus 5 --- madgraph/iolibs/export_v4.py | 50 +++++++++++++++++++++++++++++++++--- 1 file changed, 47 insertions(+), 3 deletions(-) diff --git a/madgraph/iolibs/export_v4.py b/madgraph/iolibs/export_v4.py index c1c1fb2a83..f860010a6b 100755 --- a/madgraph/iolibs/export_v4.py +++ b/madgraph/iolibs/export_v4.py @@ -2780,6 +2780,17 @@ def get_JAMP_lines(self, col_amps, JAMP_format="JAMP(%s)", AMP_format="AMP(%s)", for key in all_element: all_element[key] = complex(all_element[key]) self.jamp_orbits = None + + # Half the color flows are the reverse of the other half; where both of + # a pair weigh the same, only one is worked out and the other copied. + halving = self.get_jamp_halving( + col_amps if symmetry_source is None else symmetry_source) + if halving: + kept = set(line + 1 for line in halving['representatives']) + all_element = dict((key, value) + for key, value in all_element.items() + if key[0] in kept) + # the color basis is read from the matrix element, which is not always # what is passed here: the split order version hands over one list of # color amplitudes per order and says where they came from @@ -2913,9 +2924,16 @@ def format(frac): max_jamp = max(max_jamp, jamp) + if halving: + max_jamp = max(max_jamp, len(halving['reverse'])) for i in range(1,max_jamp+1): name = JAMP_format % i - if not jamp_res[i]: + if halving and halving['reverse'][i - 1] < i - 1: + # the reverse of a flow already worked out. Only the color sum + # skips these; everything else reads the whole array. + res_list.append(" %s = %s" % + (name, JAMP_format % (halving['reverse'][i - 1] + 1))) + elif not jamp_res[i]: res_list.append(" %s = 0d0" %(name)) else: res_list.append(" %s = %s" %(name, '+'.join(jamp_res[i]))) @@ -3227,6 +3245,20 @@ def get_jamp_folding(self, matrix_element): return {'reverse': reverse, 'sign': sign, 'representatives': representatives, 'slot': slot} + def get_jamp_halving(self, matrix_element): + """The folding, but only where it may also be used for the JAMP + definitions themselves. That needs sign +1: then both lines of a pair + weigh the same, a permutation of the lines is a permutation of the + pairs with nothing else attached, and what is left of the matrix is + still symmetric enough for the optimisation to work on. With sign -1 a + permutation may send a kept line onto its partner, which flips the + weight for that pair only.""" + + folding = self.get_jamp_folding(matrix_element) + if not folding or folding['sign'] < 0: + return None + return folding + def jamp_folded_color_matrix(self, matrix_element, reverse, sign): """The color matrix over one line per reversal pair. Summing |M|^2 over the pairs instead of over every line gives the same number, since the @@ -3340,12 +3372,21 @@ def get_jamp_symmetry(self, matrix_element, all_element): return None nb_line = len(symmetry.keys1) + # with one line per reversal pair kept, a permutation of the lines is + # read as the permutation of the pairs it induces + halving = self.get_jamp_halving(matrix_element) rowperms, actions = [], [] for induced in symmetry.generators1: - action = self.jamp_amp_permutation(columns, induced) + rowmap = induced + if halving: + kept_lines = halving['representatives'] + rowmap = list(range(nb_line)) + for line in kept_lines: + rowmap[line] = kept_lines[halving['slot'][induced[line]]] + action = self.jamp_amp_permutation(columns, rowmap) if action is None: continue - rowperms.append([0] + [induced[i] + 1 for i in range(nb_line)]) + rowperms.append([0] + [rowmap[i] + 1 for i in range(nb_line)]) actions.append(action) if not actions: return None @@ -3368,6 +3409,9 @@ def find(x): if ri != rj: parent[ri] = rj line_reps = [i for i in range(1, nb_line + 1) if find(i) == i] + if halving: + kept = set(line + 1 for line in halving['representatives']) + line_reps = [i for i in line_reps if i in kept] or sorted(kept)[:1] return {'rowperms': rowperms, 'actions': actions, 'nb_line': nb_line, 'line_reps': line_reps} From e965843b7a5ae7ce1cfb53ebe698f2f680045485 Mon Sep 17 00:00:00 2001 From: Olivier Mattelaer Date: Fri, 7 Aug 2026 00:14:14 +0200 Subject: [PATCH 15/42] Revert "work out one color flow per reversal pair and copy the other" This reverts commit 48ab222d8846e674a8ae9156229bb7d310fb6fc1. --- madgraph/iolibs/export_v4.py | 50 +++--------------------------------- 1 file changed, 3 insertions(+), 47 deletions(-) diff --git a/madgraph/iolibs/export_v4.py b/madgraph/iolibs/export_v4.py index f860010a6b..c1c1fb2a83 100755 --- a/madgraph/iolibs/export_v4.py +++ b/madgraph/iolibs/export_v4.py @@ -2780,17 +2780,6 @@ def get_JAMP_lines(self, col_amps, JAMP_format="JAMP(%s)", AMP_format="AMP(%s)", for key in all_element: all_element[key] = complex(all_element[key]) self.jamp_orbits = None - - # Half the color flows are the reverse of the other half; where both of - # a pair weigh the same, only one is worked out and the other copied. - halving = self.get_jamp_halving( - col_amps if symmetry_source is None else symmetry_source) - if halving: - kept = set(line + 1 for line in halving['representatives']) - all_element = dict((key, value) - for key, value in all_element.items() - if key[0] in kept) - # the color basis is read from the matrix element, which is not always # what is passed here: the split order version hands over one list of # color amplitudes per order and says where they came from @@ -2924,16 +2913,9 @@ def format(frac): max_jamp = max(max_jamp, jamp) - if halving: - max_jamp = max(max_jamp, len(halving['reverse'])) for i in range(1,max_jamp+1): name = JAMP_format % i - if halving and halving['reverse'][i - 1] < i - 1: - # the reverse of a flow already worked out. Only the color sum - # skips these; everything else reads the whole array. - res_list.append(" %s = %s" % - (name, JAMP_format % (halving['reverse'][i - 1] + 1))) - elif not jamp_res[i]: + if not jamp_res[i]: res_list.append(" %s = 0d0" %(name)) else: res_list.append(" %s = %s" %(name, '+'.join(jamp_res[i]))) @@ -3245,20 +3227,6 @@ def get_jamp_folding(self, matrix_element): return {'reverse': reverse, 'sign': sign, 'representatives': representatives, 'slot': slot} - def get_jamp_halving(self, matrix_element): - """The folding, but only where it may also be used for the JAMP - definitions themselves. That needs sign +1: then both lines of a pair - weigh the same, a permutation of the lines is a permutation of the - pairs with nothing else attached, and what is left of the matrix is - still symmetric enough for the optimisation to work on. With sign -1 a - permutation may send a kept line onto its partner, which flips the - weight for that pair only.""" - - folding = self.get_jamp_folding(matrix_element) - if not folding or folding['sign'] < 0: - return None - return folding - def jamp_folded_color_matrix(self, matrix_element, reverse, sign): """The color matrix over one line per reversal pair. Summing |M|^2 over the pairs instead of over every line gives the same number, since the @@ -3372,21 +3340,12 @@ def get_jamp_symmetry(self, matrix_element, all_element): return None nb_line = len(symmetry.keys1) - # with one line per reversal pair kept, a permutation of the lines is - # read as the permutation of the pairs it induces - halving = self.get_jamp_halving(matrix_element) rowperms, actions = [], [] for induced in symmetry.generators1: - rowmap = induced - if halving: - kept_lines = halving['representatives'] - rowmap = list(range(nb_line)) - for line in kept_lines: - rowmap[line] = kept_lines[halving['slot'][induced[line]]] - action = self.jamp_amp_permutation(columns, rowmap) + action = self.jamp_amp_permutation(columns, induced) if action is None: continue - rowperms.append([0] + [rowmap[i] + 1 for i in range(nb_line)]) + rowperms.append([0] + [induced[i] + 1 for i in range(nb_line)]) actions.append(action) if not actions: return None @@ -3409,9 +3368,6 @@ def find(x): if ri != rj: parent[ri] = rj line_reps = [i for i in range(1, nb_line + 1) if find(i) == i] - if halving: - kept = set(line + 1 for line in halving['representatives']) - line_reps = [i for i in line_reps if i in kept] or sorted(kept)[:1] return {'rowperms': rowperms, 'actions': actions, 'nb_line': nb_line, 'line_reps': line_reps} From 386a3ffcf5fdaae88142faee3ead432b451317bc Mon Sep 17 00:00:00 2001 From: Olivier Mattelaer Date: Fri, 7 Aug 2026 01:39:02 +0200 Subject: [PATCH 16/42] fold the color matrix only where the template sums over the pairs get_color_data_lines sits on ProcessExporterFortran and is shared by madevent, madweight and the fks exporters, but only the standalone template was taught to sum over NCOLORFOLD. So those three were handed a folded matrix and kept looping over every line: g g > g g g wrote a 78 entry triangle for the 12 pairs into CF(300) and then read all 300, the last 222 never set. jamp_fold moves to the standalone exporter, where the template agrees. The others generate exactly what they did before the folding went in. Co-Authored-By: Claude Opus 5 --- madgraph/iolibs/export_v4.py | 6 +++++- 1 file changed, 5 insertions(+), 1 deletion(-) diff --git a/madgraph/iolibs/export_v4.py b/madgraph/iolibs/export_v4.py index c1c1fb2a83..5b9be53730 100755 --- a/madgraph/iolibs/export_v4.py +++ b/madgraph/iolibs/export_v4.py @@ -230,7 +230,10 @@ class ProcessExporterFortran(VirtualExporter): # how many times the JAMP optimisation called itself, for the record myjamp_count = 0 # sum |M|^2 over one color flow per reversal pair instead of over every one - jamp_fold = True + # Folding the color matrix onto one line per reversal pair only works + # where the template sums over NCOLORFOLD. get_color_data_lines is shared + # by every fortran exporter, so this stays off unless the template agrees. + jamp_fold = False # write the JAMP definitions as one recipe per orbit of the permutations # leaving the color basis invariant, instead of one line per definition jamp_orbit = False @@ -4891,6 +4894,7 @@ class ProcessExporterFortranSA(ProcessExporterFortran): f2py_wrapper_all ="f2py_wrapper_all.inc" f2py_matrix_splitter = "f2py_splitter.py" jamp_optim = True + jamp_fold = True jamp_orbit = True default_vector_size = 0 # When True, emit per-call IAND(WF_FLAVOR_MASK/AMP_FLAVOR_MASK, From d26858b2d8fb5eb5cf096c12896cfadf6266736b Mon Sep 17 00:00:00 2001 From: Olivier Mattelaer Date: Fri, 7 Aug 2026 01:49:57 +0200 Subject: [PATCH 17/42] carry the folded color sum in the madevent templates, switched off The three madevent templates can now sum |M|^2 over one color flow per reversal pair: NCOLORFOLD, a JFOLD gathered per split order, and the sum reading %(color_fold_array)s, which is JAMP when there is no folding so the generated code is unchanged in that case -- g g > g g and g g > g g g both come back at 4.454e+08 and 3.677e+07 pb, the numbers from before the patch. It is off because the folded numbers are wrong: 2.672e+09 against 4.454e+08 on g g > g g and 3.971e+09 against 3.677e+07 on g g > g g g. Both compile and run, so this is not a missing declaration, and matrix1_optim.f -- what actually runs under helicity recycling -- is rewritten by hel_recycle out of the template, which is where I would look first: the gather sits inside the helicity loop and reads JAMP(COLREP(ICF),ICFSO), and the recycler rewrites JAMP references. get_color_data_lines folds for whichever exporter asks, so template_matrix1.f had to be taught the same sum: it comes from the group hel template and feeds the recycler. Co-Authored-By: Claude Opus 5 --- madgraph/iolibs/export_v4.py | 36 +++++++++++++++++++ .../matrix_madevent_group_v4.inc | 16 +++++---- .../matrix_madevent_group_v4_hel.inc | 16 +++++---- .../template_files/matrix_madevent_v4.inc | 16 +++++---- 4 files changed, 66 insertions(+), 18 deletions(-) diff --git a/madgraph/iolibs/export_v4.py b/madgraph/iolibs/export_v4.py index 5b9be53730..37bac4c8be 100755 --- a/madgraph/iolibs/export_v4.py +++ b/madgraph/iolibs/export_v4.py @@ -3230,6 +3230,34 @@ def get_jamp_folding(self, matrix_element): return {'reverse': reverse, 'sign': sign, 'representatives': representatives, 'slot': slot} + def get_color_fold_ampso(self, folding, ncolor): + """Template replacements for a color sum over one line per reversal + pair, where JAMP carries a second index for the split orders. Without a + folding the sum is left on JAMP itself, so nothing is copied.""" + + if not folding: + return {'ncolorfold': ncolor, + 'color_fold_decl': '', + 'color_fold_index': '', + 'color_fold_gather': '', + 'color_fold_array': 'JAMP'} + lines = [line + 1 for line in folding['representatives']] + return { + 'ncolorfold': len(lines), + 'color_fold_decl': ( + " COMPLEX*16 JFOLD(NCOLORFOLD,NAMPSO)\n" + " INTEGER COLREP(NCOLORFOLD)\n" + " INTEGER ICF, ICFSO"), + 'color_fold_index': "\n".join( + self.get_int_data_lines("COLREP", lines, var='ICF')), + 'color_fold_gather': ( + " DO ICFSO = 1, NAMPSO\n" + " DO ICF = 1, NCOLORFOLD\n" + " JFOLD(ICF,ICFSO) = JAMP(COLREP(ICF),ICFSO)\n" + " ENDDO\n" + " ENDDO"), + 'color_fold_array': 'JFOLD'} + def jamp_folded_color_matrix(self, matrix_element, reverse, sign): """The color matrix over one line per reversal pair. Summing |M|^2 over the pairs instead of over every line gives the same number, since the @@ -7065,6 +7093,9 @@ class ProcessExporterFortranME(ProcessExporterFortran): MadEvent format.""" matrix_file = "matrix_madevent_v4.inc" + # The templates carry the folded color sum, but the numbers come out + # wrong (g g > g g is 6x too large), so it stays off until that is found. + jamp_fold = False jamp_orbit = True # AMP is indexed by helicity once the matrix element is rewritten for # helicity recycling, so the definitions cannot sit at the end of it @@ -7798,6 +7829,11 @@ def write_matrix_element_v4(self, writer, matrix_element, fortran_model, # Extract ncolor ncolor = max(1, len(matrix_element.get('color_basis'))) replace_dict['ncolor'] = ncolor + # |M|^2 is summed over one color flow per reversal pair when the basis + # allows it. JAMP itself keeps every flow: jamp2 and the color flow + # selection below read all of them. + folding = self.get_jamp_folding(matrix_element) + replace_dict.update(self.get_color_fold_ampso(folding, ncolor)) # Extract color data lines color_data_lines = self.get_color_data_lines(matrix_element) diff --git a/madgraph/iolibs/template_files/matrix_madevent_group_v4.inc b/madgraph/iolibs/template_files/matrix_madevent_group_v4.inc index eb252d191d..57ee8987be 100644 --- a/madgraph/iolibs/template_files/matrix_madevent_group_v4.inc +++ b/madgraph/iolibs/template_files/matrix_madevent_group_v4.inc @@ -296,8 +296,9 @@ C include 'genps.inc' include 'nexternal.inc' include 'maxamps.inc' - INTEGER NWAVEFUNCS, NCOLOR + INTEGER NWAVEFUNCS, NCOLOR, NCOLORFOLD PARAMETER (NWAVEFUNCS=%(nwavefuncs)d, NCOLOR=%(ncolor)d) + PARAMETER (NCOLORFOLD=%(ncolorfold)d) REAL*8 ZERO PARAMETER (ZERO=0D0) COMPLEX*16 IMAG1 @@ -324,9 +325,10 @@ C COMPLEX*16 ZTEMP %(jamp_tmp_decl)s %(jamp_decl)s - INTEGER CF(NCOLOR*(NCOLOR+1)/2) + INTEGER CF(NCOLORFOLD*(NCOLORFOLD+1)/2) INTEGER DENOM, CF_INDEX COMPLEX*16 AMP(NGRAPHS), JAMP(NCOLOR,NAMPSO) +%(color_fold_decl)s type(aloha) W(NWAVEFUNCS) C Needed for v4 models COMPLEX*16 DUM0,DUM1 @@ -366,6 +368,7 @@ C C COLOR DATA C %(color_data_lines)s +%(color_fold_index)s C ---------- C BEGIN CODE C ---------- @@ -395,18 +398,19 @@ JAMP(:,:) = (0d0,0d0) ENDDO endif +%(color_fold_gather)s MATRIX%(proc_id)s = 0.D0 DO M = 1, NAMPSO CF_INDEX = 0 - DO I = 1, NCOLOR + DO I = 1, NCOLORFOLD ZTEMP = (0.D0,0.D0) - DO J = I, NCOLOR + DO J = I, NCOLORFOLD CF_INDEX = CF_INDEX + 1 - ZTEMP = ZTEMP + CF(CF_INDEX)*JAMP(J,M) + ZTEMP = ZTEMP + CF(CF_INDEX)*%(color_fold_array)s(J,M) ENDDO DO N = 1, NAMPSO %(select_configs_if)s - MATRIX%(proc_id)s = MATRIX%(proc_id)s + ZTEMP*DCONJG(JAMP(I,N)) + MATRIX%(proc_id)s = MATRIX%(proc_id)s + ZTEMP*DCONJG(%(color_fold_array)s(I,N)) %(select_configs_endif)s ENDDO ENDDO diff --git a/madgraph/iolibs/template_files/matrix_madevent_group_v4_hel.inc b/madgraph/iolibs/template_files/matrix_madevent_group_v4_hel.inc index 5cdb43bbe1..76251d5fd3 100644 --- a/madgraph/iolibs/template_files/matrix_madevent_group_v4_hel.inc +++ b/madgraph/iolibs/template_files/matrix_madevent_group_v4_hel.inc @@ -207,8 +207,9 @@ C include 'genps.inc' include 'nexternal.inc' include 'maxamps.inc' - INTEGER NWAVEFUNCS, NCOLOR + INTEGER NWAVEFUNCS, NCOLOR, NCOLORFOLD PARAMETER (NWAVEFUNCS=${nwavefuncs}, NCOLOR=%(ncolor)d) + PARAMETER (NCOLORFOLD=%(ncolorfold)d) REAL*8 ZERO PARAMETER (ZERO=0D0) COMPLEX*16 IMAG1 @@ -238,9 +239,10 @@ C %(jamp_tmp_decl)s %(jamp_decl)s COMPLEX*16 TMP(%(wavefunctionsize)d) - INTEGER CF(NCOLOR*(NCOLOR+1)) + INTEGER CF(NCOLORFOLD*(NCOLORFOLD+1)) INTEGER DENOM, CF_INDEX COMPLEX*16 AMP(NCOMB,NGRAPHS), JAMP(NCOLOR,NAMPSO) +%(color_fold_decl)s type(aloha) W(NWAVEFUNCS) C Needed for v4 models COMPLEX*16 DUM0,DUM1 @@ -274,6 +276,7 @@ C C COLOR DATA C %(color_data_lines)s +%(color_fold_index)s C ---------- C BEGIN CODE C ---------- @@ -290,18 +293,19 @@ ${helas_calls} JAMP(:,:) = (0d0,0d0) DO K = 1, NCOMB ${jamp_lines} +%(color_fold_gather)s TS(K) = 0.D0 DO M = 1, NAMPSO CF_INDEX = 0 - DO I = 1, NCOLOR + DO I = 1, NCOLORFOLD ZTEMP = (0.D0,0.D0) - DO J = I, NCOLOR + DO J = I, NCOLORFOLD CF_INDEX = CF_INDEX + 1 - ZTEMP = ZTEMP + CF(CF_INDEX)*JAMP(J,M) + ZTEMP = ZTEMP + CF(CF_INDEX)*%(color_fold_array)s(J,M) ENDDO ! J DO N = 1, NAMPSO %(select_configs_if)s - TS(K) = TS(K) + REAL(ZTEMP*DCONJG(JAMP(I,N))) + TS(K) = TS(K) + REAL(ZTEMP*DCONJG(%(color_fold_array)s(I,N))) %(select_configs_endif)s ENDDO ! N ENDDO ! I diff --git a/madgraph/iolibs/template_files/matrix_madevent_v4.inc b/madgraph/iolibs/template_files/matrix_madevent_v4.inc index a78cfbe917..24759f00b8 100644 --- a/madgraph/iolibs/template_files/matrix_madevent_v4.inc +++ b/madgraph/iolibs/template_files/matrix_madevent_v4.inc @@ -246,8 +246,9 @@ C include 'genps.inc' include 'nexternal.inc' include 'maxamps.inc' - INTEGER NWAVEFUNCS, NCOLOR + INTEGER NWAVEFUNCS, NCOLOR, NCOLORFOLD PARAMETER (NWAVEFUNCS=%(nwavefuncs)d, NCOLOR=%(ncolor)d) + PARAMETER (NCOLORFOLD=%(ncolorfold)d) REAL*8 ZERO PARAMETER (ZERO=0D0) COMPLEX*16 IMAG1 @@ -274,9 +275,10 @@ C COMPLEX*16 ZTEMP %(jamp_tmp_decl)s %(jamp_decl)s - INTEGER CF(NCOLOR*(NCOLOR+1)) + INTEGER CF(NCOLORFOLD*(NCOLORFOLD+1)) INTEGER CF_INDEX,DENOM COMPLEX*16 AMP(NGRAPHS), JAMP(NCOLOR,NAMPSO) +%(color_fold_decl)s type(aloha) W(NWAVEFUNCS) C Needed for v4 models COMPLEX*16 DUM0,DUM1 @@ -304,6 +306,7 @@ C C COLOR DATA C %(color_data_lines)s +%(color_fold_index)s C ---------- C BEGIN CODE C ---------- @@ -319,18 +322,19 @@ AMP(:) = (0d0,0d0) %(helas_calls)s %(jamp_lines)s +%(color_fold_gather)s MATRIX%(proc_id)s = 0.D0 DO M = 1, NAMPSO CF_INDEX = 0 - DO I = 1, NCOLOR + DO I = 1, NCOLORFOLD ZTEMP = (0.D0,0.D0) - DO J = I, NCOLOR + DO J = I, NCOLORFOLD CF_INDEX = CF_INDEX +1 - ZTEMP = ZTEMP + CF(CF_INDEX)*JAMP(J,M) + ZTEMP = ZTEMP + CF(CF_INDEX)*%(color_fold_array)s(J,M) ENDDO DO N = 1, NAMPSO IF (CHOSEN_SO_CONFIGS(SQSOINDEX%(proc_id)s(M,N))) THEN - MATRIX%(proc_id)s = MATRIX%(proc_id)s + ZTEMP*DCONJG(JAMP(I,N)) + MATRIX%(proc_id)s = MATRIX%(proc_id)s + ZTEMP*DCONJG(%(color_fold_array)s(I,N)) ENDIF ENDDO ENDDO From 10a1dad56cbf1e19d9b24b5416dc31c7c43a0dfd Mon Sep 17 00:00:00 2001 From: Olivier Mattelaer Date: Fri, 7 Aug 2026 02:06:25 +0200 Subject: [PATCH 18/42] gather the folded color flows without opening a block hel_recycle scrapes the lines between the jamp block and the color sum into amp2_lines, starting at the first ENDDO after the jamp lines and stopping at a six space indented "DO I = 1, NCOLOR". The gather loop put an ENDDO in front of that start, and the folded sum sits two levels in, so the whole sum was scraped and emitted a second time inside the sde_strat branch -- without its DENOM line, which that scrape drops on purpose. The second copy overwrote TS(K), so |M|^2 came out DENOM times too large: 6x on g g > g g, 108x on g g > g g g, each exactly the DENOM of that process. A vector subscript gathers in one statement and leaves the scrape alone. The recycler is unchanged: it was reading what it was written to read. g g > g g and g g > g g g both come back at 4.454e+08 and 3.677e+07 pb with the folding on and the recycler doing its work, the same to every digit quoted, uncertainties included, as the unfolded runs. TS(K) is zeroed once and keeps its DENOM. Co-Authored-By: Claude Opus 5 --- madgraph/iolibs/export_v4.py | 13 +++---------- 1 file changed, 3 insertions(+), 10 deletions(-) diff --git a/madgraph/iolibs/export_v4.py b/madgraph/iolibs/export_v4.py index 37bac4c8be..bf9a8bd231 100755 --- a/madgraph/iolibs/export_v4.py +++ b/madgraph/iolibs/export_v4.py @@ -3247,15 +3247,10 @@ def get_color_fold_ampso(self, folding, ncolor): 'color_fold_decl': ( " COMPLEX*16 JFOLD(NCOLORFOLD,NAMPSO)\n" " INTEGER COLREP(NCOLORFOLD)\n" - " INTEGER ICF, ICFSO"), + " INTEGER ICF"), 'color_fold_index': "\n".join( self.get_int_data_lines("COLREP", lines, var='ICF')), - 'color_fold_gather': ( - " DO ICFSO = 1, NAMPSO\n" - " DO ICF = 1, NCOLORFOLD\n" - " JFOLD(ICF,ICFSO) = JAMP(COLREP(ICF),ICFSO)\n" - " ENDDO\n" - " ENDDO"), + 'color_fold_gather': " JFOLD(:,:) = JAMP(COLREP(:),:)", 'color_fold_array': 'JFOLD'} def jamp_folded_color_matrix(self, matrix_element, reverse, sign): @@ -7093,9 +7088,7 @@ class ProcessExporterFortranME(ProcessExporterFortran): MadEvent format.""" matrix_file = "matrix_madevent_v4.inc" - # The templates carry the folded color sum, but the numbers come out - # wrong (g g > g g is 6x too large), so it stays off until that is found. - jamp_fold = False + jamp_fold = True jamp_orbit = True # AMP is indexed by helicity once the matrix element is rewritten for # helicity recycling, so the definitions cannot sit at the end of it From 9fe3443338ce8e9752e09036b78110d4e7af2e22 Mon Sep 17 00:00:00 2001 From: Olivier Mattelaer Date: Fri, 7 Aug 2026 02:16:46 +0200 Subject: [PATCH 19/42] sum |M|^2 with real arithmetic The color matrix is real and symmetric and MATRIX is real, so the imaginary part of the sum is thrown away as it stands. Splitting JAMP into its two parts computes what is kept and nothing else: Re( sum_j C_ij JAMP_j conj(JAMP_i) ) = (sum_j C_ij ReJ_j) ReJ_i + (sum_j C_ij ImJ_j) ImJ_i term by term, so it needs no cancellation argument and holds for every process, not only the ones whose coefficients are real. The gather fills two REAL*8 arrays instead of one COMPLEX*16, and the four accumulators become eight. g g > g g g, g g > g g g g and g g > t t~ g g are unchanged to the last digit. u u~ > u u~ g g moves in the last digit, the same 1 ulp it already carried before this. Co-Authored-By: Claude Opus 5 --- madgraph/iolibs/export_v4.py | 6 ++- .../template_files/matrix_standalone_v4.inc | 38 ++++++++++++------- 2 files changed, 29 insertions(+), 15 deletions(-) diff --git a/madgraph/iolibs/export_v4.py b/madgraph/iolibs/export_v4.py index bf9a8bd231..268de26c72 100755 --- a/madgraph/iolibs/export_v4.py +++ b/madgraph/iolibs/export_v4.py @@ -5827,10 +5827,12 @@ def write_matrix_element_v4(self, writer, matrix_element, fortran_model, var='ICF')) if folding else '' replace_dict['color_fold_gather'] = ( " DO ICF = 1, NCOLORFOLD\n" - " JFOLD(ICF) = JAMP(COLREP(ICF))\n" + " JFR(ICF) = DBLE(JAMP(COLREP(ICF)))\n" + " JFI(ICF) = DIMAG(JAMP(COLREP(ICF)))\n" " ENDDO" if folding else " DO ICF = 1, NCOLOR\n" - " JFOLD(ICF) = JAMP(ICF)\n" + " JFR(ICF) = DBLE(JAMP(ICF))\n" + " JFI(ICF) = DIMAG(JAMP(ICF))\n" " ENDDO") if not folding: replace_dict['color_fold_decl'] = " INTEGER ICF" diff --git a/madgraph/iolibs/template_files/matrix_standalone_v4.inc b/madgraph/iolibs/template_files/matrix_standalone_v4.inc index db5f614509..553fd96547 100644 --- a/madgraph/iolibs/template_files/matrix_standalone_v4.inc +++ b/madgraph/iolibs/template_files/matrix_standalone_v4.inc @@ -374,14 +374,16 @@ C C LOCAL VARIABLES C INTEGER I,J,NJ,NB - COMPLEX*16 ZTEMP,Z1,Z2,Z3,Z4 + REAL*8 ZR,ZI,ZR1,ZR2,ZR3,ZR4,ZI1,ZI2,ZI3,ZI4 INTEGER CF_INDEX INTEGER %(proc_prefix)sCF(NCOLORFOLD*(NCOLORFOLD+1)/2) INTEGER %(proc_prefix)sDENOM common /%(proc_prefix)scolor_matrix/ %(proc_prefix)sCF,%(proc_prefix)sDENOM COMPLEX*16 JAMP(NCOLOR) - COMPLEX*16 JFOLD(NCOLORFOLD) +C The color matrix is real and symmetric and MATRIX is real, so the +C sum splits into two real quadratic forms, one on each part of JAMP. + REAL*8 JFR(NCOLORFOLD), JFI(NCOLORFOLD) %(color_fold_decl)s %(color_fold_index)s %(jamp_tmp_decl)s @@ -405,24 +407,34 @@ C adder, and that latency is what the loop spends its C time on. No compiler does this by itself, since it C changes the order the terms are summed in. DO I = 1, NCOLORFOLD - Z1 = (0.D0,0.D0) - Z2 = (0.D0,0.D0) - Z3 = (0.D0,0.D0) - Z4 = (0.D0,0.D0) + ZR1 = 0.D0 + ZR2 = 0.D0 + ZR3 = 0.D0 + ZR4 = 0.D0 + ZI1 = 0.D0 + ZI2 = 0.D0 + ZI3 = 0.D0 + ZI4 = 0.D0 NJ = NCOLORFOLD - I + 1 NB = (NJ/4)*4 DO J = 0, NB-4, 4 - Z1 = Z1 + %(proc_prefix)sCF(CF_INDEX+J+1)*JFOLD(I+J) - Z2 = Z2 + %(proc_prefix)sCF(CF_INDEX+J+2)*JFOLD(I+J+1) - Z3 = Z3 + %(proc_prefix)sCF(CF_INDEX+J+3)*JFOLD(I+J+2) - Z4 = Z4 + %(proc_prefix)sCF(CF_INDEX+J+4)*JFOLD(I+J+3) + ZR1 = ZR1 + %(proc_prefix)sCF(CF_INDEX+J+1)*JFR(I+J) + ZR2 = ZR2 + %(proc_prefix)sCF(CF_INDEX+J+2)*JFR(I+J+1) + ZR3 = ZR3 + %(proc_prefix)sCF(CF_INDEX+J+3)*JFR(I+J+2) + ZR4 = ZR4 + %(proc_prefix)sCF(CF_INDEX+J+4)*JFR(I+J+3) + ZI1 = ZI1 + %(proc_prefix)sCF(CF_INDEX+J+1)*JFI(I+J) + ZI2 = ZI2 + %(proc_prefix)sCF(CF_INDEX+J+2)*JFI(I+J+1) + ZI3 = ZI3 + %(proc_prefix)sCF(CF_INDEX+J+3)*JFI(I+J+2) + ZI4 = ZI4 + %(proc_prefix)sCF(CF_INDEX+J+4)*JFI(I+J+3) ENDDO - ZTEMP = (Z1+Z2)+(Z3+Z4) + ZR = (ZR1+ZR2)+(ZR3+ZR4) + ZI = (ZI1+ZI2)+(ZI3+ZI4) DO J = NB, NJ-1 - ZTEMP = ZTEMP + %(proc_prefix)sCF(CF_INDEX+J+1)*JFOLD(I+J) + ZR = ZR + %(proc_prefix)sCF(CF_INDEX+J+1)*JFR(I+J) + ZI = ZI + %(proc_prefix)sCF(CF_INDEX+J+1)*JFI(I+J) ENDDO CF_INDEX = CF_INDEX + NJ - MATRIX = MATRIX+ZTEMP*DCONJG(JFOLD(I))/%(proc_prefix)sDENOM + MATRIX = MATRIX+(ZR*JFR(I)+ZI*JFI(I))/%(proc_prefix)sDENOM ENDDO END From ee8a18b731972e037858465f0893b494e38ee9e7 Mon Sep 17 00:00:00 2001 From: Olivier Mattelaer Date: Fri, 7 Aug 2026 02:22:21 +0200 Subject: [PATCH 20/42] refresh the madevent references for the folded color sum Both are the case where the color basis has no reversal symmetry, so the folding declines and the only change is the name of the loop bound: NCOLORFOLD equals NCOLOR, the gather is empty and the sum still reads JAMP. testIO_export_matrix_element_v4_standalone still fails, from 9fe344333 rather than from here: that commit moved the standalone sum to real arithmetic and left the reference on the complex form. Co-Authored-By: Claude Opus 5 --- .../matrix1.f | 12 ++++++++---- .../matrix.f | 12 ++++++++---- 2 files changed, 16 insertions(+), 8 deletions(-) diff --git a/tests/input_files/IOTestsComparison/IOExportV4IOTest/export_matrix_element_v4_madevent_group/matrix1.f b/tests/input_files/IOTestsComparison/IOExportV4IOTest/export_matrix_element_v4_madevent_group/matrix1.f index 25bc11f65a..d5a1ff7ade 100644 --- a/tests/input_files/IOTestsComparison/IOExportV4IOTest/export_matrix_element_v4_madevent_group/matrix1.f +++ b/tests/input_files/IOTestsComparison/IOExportV4IOTest/export_matrix_element_v4_madevent_group/matrix1.f @@ -372,8 +372,9 @@ REAL*8 FUNCTION MATRIX1(P,NHEL,IFLAV, IHEL,AMP2, JAMP2, IVEC) INCLUDE 'genps.inc' INCLUDE 'nexternal.inc' INCLUDE 'maxamps.inc' - INTEGER NWAVEFUNCS, NCOLOR + INTEGER NWAVEFUNCS, NCOLOR, NCOLORFOLD PARAMETER (NWAVEFUNCS=5, NCOLOR=2) + PARAMETER (NCOLORFOLD=2) REAL*8 ZERO PARAMETER (ZERO=0D0) COMPLEX*16 IMAG1 @@ -400,9 +401,10 @@ REAL*8 FUNCTION MATRIX1(P,NHEL,IFLAV, IHEL,AMP2, JAMP2, IVEC) COMPLEX*16 ZTEMP COMPLEX*16 TMP_JAMP(0) - INTEGER CF(NCOLOR*(NCOLOR+1)/2) + INTEGER CF(NCOLORFOLD*(NCOLORFOLD+1)/2) INTEGER DENOM, CF_INDEX COMPLEX*16 AMP(NGRAPHS), JAMP(NCOLOR,NAMPSO) + TYPE(ALOHA) W(NWAVEFUNCS) C Needed for v4 models COMPLEX*16 DUM0,DUM1 @@ -450,6 +452,7 @@ REAL*8 FUNCTION MATRIX1(P,NHEL,IFLAV, IHEL,AMP2, JAMP2, IVEC) C 1 T(2,1) T(3,4) DATA (CF(I),I= 3, 3) /9/ C 1 T(2,4) T(3,1) + C ---------- C BEGIN CODE C ---------- @@ -512,12 +515,13 @@ REAL*8 FUNCTION MATRIX1(P,NHEL,IFLAV, IHEL,AMP2, JAMP2, IVEC) ENDDO ENDIF + MATRIX1 = 0.D0 DO M = 1, NAMPSO CF_INDEX = 0 - DO I = 1, NCOLOR + DO I = 1, NCOLORFOLD ZTEMP = (0.D0,0.D0) - DO J = I, NCOLOR + DO J = I, NCOLORFOLD CF_INDEX = CF_INDEX + 1 ZTEMP = ZTEMP + CF(CF_INDEX)*JAMP(J,M) ENDDO diff --git a/tests/input_files/IOTestsComparison/IOExportV4IOTest/export_matrix_element_v4_madevent_nogroup/matrix.f b/tests/input_files/IOTestsComparison/IOExportV4IOTest/export_matrix_element_v4_madevent_nogroup/matrix.f index 3ee6679790..d339029b1d 100644 --- a/tests/input_files/IOTestsComparison/IOExportV4IOTest/export_matrix_element_v4_madevent_nogroup/matrix.f +++ b/tests/input_files/IOTestsComparison/IOExportV4IOTest/export_matrix_element_v4_madevent_nogroup/matrix.f @@ -307,8 +307,9 @@ REAL*8 FUNCTION MATRIX(P,NHEL,IFLAV, IVEC) INCLUDE 'genps.inc' INCLUDE 'nexternal.inc' INCLUDE 'maxamps.inc' - INTEGER NWAVEFUNCS, NCOLOR + INTEGER NWAVEFUNCS, NCOLOR, NCOLORFOLD PARAMETER (NWAVEFUNCS=5, NCOLOR=2) + PARAMETER (NCOLORFOLD=2) REAL*8 ZERO PARAMETER (ZERO=0D0) COMPLEX*16 IMAG1 @@ -335,9 +336,10 @@ REAL*8 FUNCTION MATRIX(P,NHEL,IFLAV, IVEC) COMPLEX*16 ZTEMP COMPLEX*16 TMP_JAMP(0) - INTEGER CF(NCOLOR*(NCOLOR+1)) + INTEGER CF(NCOLORFOLD*(NCOLORFOLD+1)) INTEGER CF_INDEX,DENOM COMPLEX*16 AMP(NGRAPHS), JAMP(NCOLOR,NAMPSO) + TYPE(ALOHA) W(NWAVEFUNCS) C Needed for v4 models COMPLEX*16 DUM0,DUM1 @@ -371,6 +373,7 @@ REAL*8 FUNCTION MATRIX(P,NHEL,IFLAV, IVEC) C 1 T(3,4,2,1) DATA (CF(I),I= 3, 3) /16/ C 1 T(4,3,2,1) + C ---------- C BEGIN CODE C ---------- @@ -403,12 +406,13 @@ REAL*8 FUNCTION MATRIX(P,NHEL,IFLAV, IVEC) JAMP(2,1) = ((0.000000000000000D+00,1.000000000000000D+00)) $ *AMP(1)+AMP(2) + MATRIX = 0.D0 DO M = 1, NAMPSO CF_INDEX = 0 - DO I = 1, NCOLOR + DO I = 1, NCOLORFOLD ZTEMP = (0.D0,0.D0) - DO J = I, NCOLOR + DO J = I, NCOLORFOLD CF_INDEX = CF_INDEX +1 ZTEMP = ZTEMP + CF(CF_INDEX)*JAMP(J,M) ENDDO From b7291d7c2001379f09ab3ea14eb1a589cb147681 Mon Sep 17 00:00:00 2001 From: Olivier Mattelaer Date: Fri, 7 Aug 2026 02:32:49 +0200 Subject: [PATCH 21/42] Revert "sum |M|^2 with real arithmetic" This reverts commit 9fe3443338ce8e9752e09036b78110d4e7af2e22. --- madgraph/iolibs/export_v4.py | 6 +-- .../template_files/matrix_standalone_v4.inc | 38 +++++++------------ 2 files changed, 15 insertions(+), 29 deletions(-) diff --git a/madgraph/iolibs/export_v4.py b/madgraph/iolibs/export_v4.py index 268de26c72..bf9a8bd231 100755 --- a/madgraph/iolibs/export_v4.py +++ b/madgraph/iolibs/export_v4.py @@ -5827,12 +5827,10 @@ def write_matrix_element_v4(self, writer, matrix_element, fortran_model, var='ICF')) if folding else '' replace_dict['color_fold_gather'] = ( " DO ICF = 1, NCOLORFOLD\n" - " JFR(ICF) = DBLE(JAMP(COLREP(ICF)))\n" - " JFI(ICF) = DIMAG(JAMP(COLREP(ICF)))\n" + " JFOLD(ICF) = JAMP(COLREP(ICF))\n" " ENDDO" if folding else " DO ICF = 1, NCOLOR\n" - " JFR(ICF) = DBLE(JAMP(ICF))\n" - " JFI(ICF) = DIMAG(JAMP(ICF))\n" + " JFOLD(ICF) = JAMP(ICF)\n" " ENDDO") if not folding: replace_dict['color_fold_decl'] = " INTEGER ICF" diff --git a/madgraph/iolibs/template_files/matrix_standalone_v4.inc b/madgraph/iolibs/template_files/matrix_standalone_v4.inc index 553fd96547..db5f614509 100644 --- a/madgraph/iolibs/template_files/matrix_standalone_v4.inc +++ b/madgraph/iolibs/template_files/matrix_standalone_v4.inc @@ -374,16 +374,14 @@ C C LOCAL VARIABLES C INTEGER I,J,NJ,NB - REAL*8 ZR,ZI,ZR1,ZR2,ZR3,ZR4,ZI1,ZI2,ZI3,ZI4 + COMPLEX*16 ZTEMP,Z1,Z2,Z3,Z4 INTEGER CF_INDEX INTEGER %(proc_prefix)sCF(NCOLORFOLD*(NCOLORFOLD+1)/2) INTEGER %(proc_prefix)sDENOM common /%(proc_prefix)scolor_matrix/ %(proc_prefix)sCF,%(proc_prefix)sDENOM COMPLEX*16 JAMP(NCOLOR) -C The color matrix is real and symmetric and MATRIX is real, so the -C sum splits into two real quadratic forms, one on each part of JAMP. - REAL*8 JFR(NCOLORFOLD), JFI(NCOLORFOLD) + COMPLEX*16 JFOLD(NCOLORFOLD) %(color_fold_decl)s %(color_fold_index)s %(jamp_tmp_decl)s @@ -407,34 +405,24 @@ C adder, and that latency is what the loop spends its C time on. No compiler does this by itself, since it C changes the order the terms are summed in. DO I = 1, NCOLORFOLD - ZR1 = 0.D0 - ZR2 = 0.D0 - ZR3 = 0.D0 - ZR4 = 0.D0 - ZI1 = 0.D0 - ZI2 = 0.D0 - ZI3 = 0.D0 - ZI4 = 0.D0 + Z1 = (0.D0,0.D0) + Z2 = (0.D0,0.D0) + Z3 = (0.D0,0.D0) + Z4 = (0.D0,0.D0) NJ = NCOLORFOLD - I + 1 NB = (NJ/4)*4 DO J = 0, NB-4, 4 - ZR1 = ZR1 + %(proc_prefix)sCF(CF_INDEX+J+1)*JFR(I+J) - ZR2 = ZR2 + %(proc_prefix)sCF(CF_INDEX+J+2)*JFR(I+J+1) - ZR3 = ZR3 + %(proc_prefix)sCF(CF_INDEX+J+3)*JFR(I+J+2) - ZR4 = ZR4 + %(proc_prefix)sCF(CF_INDEX+J+4)*JFR(I+J+3) - ZI1 = ZI1 + %(proc_prefix)sCF(CF_INDEX+J+1)*JFI(I+J) - ZI2 = ZI2 + %(proc_prefix)sCF(CF_INDEX+J+2)*JFI(I+J+1) - ZI3 = ZI3 + %(proc_prefix)sCF(CF_INDEX+J+3)*JFI(I+J+2) - ZI4 = ZI4 + %(proc_prefix)sCF(CF_INDEX+J+4)*JFI(I+J+3) + Z1 = Z1 + %(proc_prefix)sCF(CF_INDEX+J+1)*JFOLD(I+J) + Z2 = Z2 + %(proc_prefix)sCF(CF_INDEX+J+2)*JFOLD(I+J+1) + Z3 = Z3 + %(proc_prefix)sCF(CF_INDEX+J+3)*JFOLD(I+J+2) + Z4 = Z4 + %(proc_prefix)sCF(CF_INDEX+J+4)*JFOLD(I+J+3) ENDDO - ZR = (ZR1+ZR2)+(ZR3+ZR4) - ZI = (ZI1+ZI2)+(ZI3+ZI4) + ZTEMP = (Z1+Z2)+(Z3+Z4) DO J = NB, NJ-1 - ZR = ZR + %(proc_prefix)sCF(CF_INDEX+J+1)*JFR(I+J) - ZI = ZI + %(proc_prefix)sCF(CF_INDEX+J+1)*JFI(I+J) + ZTEMP = ZTEMP + %(proc_prefix)sCF(CF_INDEX+J+1)*JFOLD(I+J) ENDDO CF_INDEX = CF_INDEX + NJ - MATRIX = MATRIX+(ZR*JFR(I)+ZI*JFI(I))/%(proc_prefix)sDENOM + MATRIX = MATRIX+ZTEMP*DCONJG(JFOLD(I))/%(proc_prefix)sDENOM ENDDO END From 58879ec55811f36d8458dcd6f0e8588c3a940d55 Mon Sep 17 00:00:00 2001 From: Olivier Mattelaer Date: Fri, 7 Aug 2026 08:08:56 +0200 Subject: [PATCH 22/42] walk whole numbers when every coefficient shares its power of i A pure gluon process picks up one factor of i per f^abc, the same for every term, so the coefficient matrix is either real throughout or imaginary throughout. Dividing that factor out leaves whole numbers for optimise_jamp to compare and hash, with nothing widened to complex, and it goes back onto the JAMP coefficients afterwards -- the definitions hold ratios, which it cancels out of. g g > 6g generates a byte identical matrix.f either way. It is worth less than it looks: the color flow step goes from 136s to 131s, under 4%, on the process where it costs the most. A quark line mixes the two phases and gets the old path. Co-Authored-By: Claude Opus 5 --- madgraph/iolibs/export_v4.py | 57 ++++++++++++++++++++++++++++++++++-- 1 file changed, 55 insertions(+), 2 deletions(-) diff --git a/madgraph/iolibs/export_v4.py b/madgraph/iolibs/export_v4.py index bf9a8bd231..cf8d11040d 100755 --- a/madgraph/iolibs/export_v4.py +++ b/madgraph/iolibs/export_v4.py @@ -234,6 +234,7 @@ class ProcessExporterFortran(VirtualExporter): # where the template sums over NCOLORFOLD. get_color_data_lines is shared # by every fortran exporter, so this stays off unless the template agrees. jamp_fold = False + jamp_integer_walk = True # write the JAMP definitions as one recipe per orbit of the permutations # leaving the color basis invariant, instead of one line per definition jamp_orbit = False @@ -2780,8 +2781,31 @@ def get_JAMP_lines(self, col_amps, JAMP_format="JAMP(%s)", AMP_format="AMP(%s)", res_list = [] self.myjamp_count = 0 - for key in all_element: - all_element[key] = complex(all_element[key]) + # With one power of i shared by every coefficient, dividing it out + # leaves whole numbers to walk over -- they compare and hash exactly, + # and nothing has to be widened to complex. The phase goes back onto + # the JAMP coefficients afterwards, so the lines written are the same. + phase = self.jamp_global_phase(all_element) \ + if self.jamp_integer_walk else None + integral = False + if phase is not None: + whole = {} + for key, value in all_element.items(): + number = value / phase if phase != 1 else value + if isinstance(number, complex): + number = number.real + number = fractions.Fraction(number).limit_denominator(10**9) + if number.denominator != 1: + break + whole[key] = int(number) + else: + all_element.clear() + all_element.update(whole) + integral = True + if not integral: + phase = None + for key in all_element: + all_element[key] = complex(all_element[key]) self.jamp_orbits = None # the color basis is read from the matrix element, which is not always # what is passed here: the split order version hands over one list of @@ -2790,6 +2814,11 @@ def get_JAMP_lines(self, col_amps, JAMP_format="JAMP(%s)", AMP_format="AMP(%s)", col_amps if symmetry_source is None else symmetry_source, all_element) if orbit and self.jamp_orbit else None new_mat, defs = self.optimise_jamp(all_element, symmetry=symmetry) + if phase is not None and phase != 1: + # the definitions hold ratios, which the phase cancels out of; only + # the coefficients on the JAMP lines carry it + for key in new_mat: + new_mat[key] = new_mat[key] * phase if start_time: logger.info("Color-Flow passed to %s term in %ss. Introduce %i contraction", len(new_mat), int(time.time()-start_time), len(defs)) @@ -3205,6 +3234,30 @@ def get_jamp_reflection(self, matrix_element): # get_jamp_folding). color_fold_max_written = 300000 + @staticmethod + def jamp_global_phase(all_element): + """The power of i every coefficient carries, when they all carry the + same one. A pure gluon process picks up one factor of i per f^abc, the + same for every term, so the whole matrix is real or wholly imaginary; + a quark line mixes the two and there is nothing to take out.""" + + phase = None + for value in all_element.values(): + if not value: + continue + number = complex(value) + if number.imag == 0: + here = 1 + elif number.real == 0: + here = 1j + else: + return None + if phase is None: + phase = here + elif phase != here: + return None + return phase + def get_jamp_folding(self, matrix_element): """Whether to sum |M|^2 over one line per reversal pair, and the (reverse, sign, representatives, slot) that goes with it. From 990cd7f69e0e8244cd9b22f14e6f35d92ee4fc21 Mon Sep 17 00:00:00 2001 From: Olivier Mattelaer Date: Fri, 7 Aug 2026 09:35:13 +0200 Subject: [PATCH 23/42] sum the color for every helicity at once, through BLAS DSYMM takes the whole helicity set as one right hand side instead of walking the color matrix once per helicity. On a synthetic g g > 6g sized problem the kernel goes from 2.485 to 0.069 ms per helicity, and that is with DSYMM doing about twice the arithmetic: it reads the full symmetric matrix where the scalar loop reads a packed triangle. DSYMM is real, so the two parts of JAMP go through separately -- the color matrix is real and symmetric, so nothing else is needed. The full matrix is built once from the written triangle, whose off diagonal is doubled because the scalar sum walks it once; halving it there is exact, the entries are integers. Only in the steady state: USERHEL unset, the good helicities settled and no polarization selection. Discovery and filtering keep the scalar path, which also stays as the fallback whenever BLAS is not taken. blas defaults to None, meaning take it when DSYMM links and the folded matrix has at least blas_min_ncolor rows -- below that the call is not worth setting up. g g > 5g agrees with the scalar build to 1.4 ulp, which is what reassociat- ing the sum costs. Co-Authored-By: Claude Opus 5 --- madgraph/iolibs/export_v4.py | 185 +++++++++++++++++- .../iolibs/template_files/makefile_sa_f_sp | 3 +- .../matrix_madevent_group_v4.inc | 5 +- .../matrix_madevent_group_v4_hel.inc | 5 +- .../template_files/matrix_madevent_v4.inc | 5 +- .../template_files/matrix_standalone_v4.inc | 5 +- 6 files changed, 201 insertions(+), 7 deletions(-) diff --git a/madgraph/iolibs/export_v4.py b/madgraph/iolibs/export_v4.py index cf8d11040d..7821529218 100755 --- a/madgraph/iolibs/export_v4.py +++ b/madgraph/iolibs/export_v4.py @@ -235,6 +235,11 @@ class ProcessExporterFortran(VirtualExporter): # by every fortran exporter, so this stays off unless the template agrees. jamp_fold = False jamp_integer_walk = True + # BLAS-3 for the color sum: all helicities at once as one right hand side. + # None means take it when the library is there and the process is big + # enough for it to pay. + blas = None + blas_min_ncolor = 100 # write the JAMP definitions as one recipe per orbit of the permutations # leaving the color basis invariant, instead of one line per definition jamp_orbit = False @@ -3234,6 +3239,61 @@ def get_jamp_reflection(self, matrix_element): # get_jamp_folding). color_fold_max_written = 300000 + _blas_available = None + + @classmethod + def blas_is_available(cls): + """Whether a BLAS carrying DSYMM can be linked, asked once.""" + + if cls._blas_available is None: + import subprocess, tempfile, shutil + probe = (" PROGRAM P\n" + " DOUBLE PRECISION A(1,1),B(1,1),C(1,1)\n" + " A=1D0\n B=1D0\n C=0D0\n" + " CALL DSYMM('L','U',1,1,1D0,A,1,B,1,0D0,C,1)\n" + " END\n") + work = tempfile.mkdtemp() + cls._blas_available = False + cls._blas_flags = '' + try: + src = os.path.join(work, 'p.f') + open(src, 'w').write(probe) + for flags in ('-framework Accelerate', '-lblas'): + try: + out = subprocess.call( + ['gfortran', src, '-o', os.path.join(work, 'p')] + + flags.split(), + stdout=subprocess.DEVNULL, + stderr=subprocess.DEVNULL) + except OSError: + break + if out == 0: + cls._blas_available = True + cls._blas_flags = flags + break + finally: + shutil.rmtree(work, ignore_errors=True) + return cls._blas_available + + def blas_link_flags(self): + """What to link the color sum against, empty when BLAS is not taken.""" + + if self.blas is False or not self.blas_is_available(): + return '' + return self._blas_flags + + def blas_wanted(self, nfold): + """Take BLAS when asked for it, or when it is there and the color + matrix is big enough that the call is worth setting up.""" + + if self.blas is False: + return False + if not self.blas_is_available(): + return False + if self.blas is True: + return True + return nfold >= self.blas_min_ncolor + @staticmethod def jamp_global_phase(all_element): """The power of i every coefficient carries, when they all carry the @@ -3283,6 +3343,64 @@ def get_jamp_folding(self, matrix_element): return {'reverse': reverse, 'sign': sign, 'representatives': representatives, 'slot': slot} + @staticmethod + def get_blas_routine(prefix, nfold, ncomb): + """The color sum for every helicity at once. DSYMM is real, so the + two parts of JAMP go through separately; the color matrix is real and + symmetric so that is all it takes.""" + + return """ + SUBROUTINE {p}GET_MATRIX_BATCH(JR,JI,NB,ANS) + IMPLICIT NONE + INTEGER NFOLD, NCOMB + PARAMETER (NFOLD={n}) + PARAMETER (NCOMB={c}) + DOUBLE PRECISION JR(NFOLD,NCOMB), JI(NFOLD,NCOMB) + INTEGER NB + DOUBLE PRECISION ANS + INTEGER I,J,K,CFI + DOUBLE PRECISION, ALLOCATABLE, SAVE :: CFULL(:,:) + DOUBLE PRECISION, ALLOCATABLE, SAVE :: TR(:,:), TI(:,:) + LOGICAL FIRST + DATA FIRST /.TRUE./ + SAVE FIRST + INTEGER {p}CF(NFOLD*(NFOLD+1)/2) + INTEGER {p}DENOM + common /{p}color_matrix/ {p}CF,{p}DENOM + IF (FIRST) THEN + CALL {p}INIT_CF() + ALLOCATE(CFULL(NFOLD,NFOLD)) + ALLOCATE(TR(NFOLD,NCOMB)) + ALLOCATE(TI(NFOLD,NCOMB)) +C What is written out is the upper triangle with its off diagonal +C doubled, since the scalar sum walks it once. BLAS wants the whole +C matrix with each entry counted once. + CFI = 0 + DO I = 1, NFOLD + DO J = I, NFOLD + CFI = CFI + 1 + IF (I.EQ.J) THEN + CFULL(I,J) = DBLE({p}CF(CFI)) + ELSE + CFULL(I,J) = DBLE({p}CF(CFI))/2D0 + CFULL(J,I) = CFULL(I,J) + ENDIF + ENDDO + ENDDO + FIRST = .FALSE. + ENDIF + CALL DSYMM('L','U',NFOLD,NB,1D0,CFULL,NFOLD,JR,NFOLD,0D0,TR,NFOLD) + CALL DSYMM('L','U',NFOLD,NB,1D0,CFULL,NFOLD,JI,NFOLD,0D0,TI,NFOLD) + ANS = 0D0 + DO K = 1, NB + DO I = 1, NFOLD + ANS = ANS + TR(I,K)*JR(I,K) + TI(I,K)*JI(I,K) + ENDDO + ENDDO + ANS = ANS / DBLE({p}DENOM) + END +""".format(p=prefix, n=nfold, c=ncomb) + def get_color_fold_ampso(self, folding, ncolor): """Template replacements for a color sum over one line per reversal pair, where JAMP carries a second index for the split orders. Without a @@ -5033,14 +5151,17 @@ def copy_template(self, model): text = fsock.read() fsock.close() fsock = open(pjoin(self.dir_path, 'SubProcesses', 'makefileP'),'w') + text = text.replace('BLASLIBS =', 'BLASLIBS = %s' % self.blas_link_flags()) text = text.replace('LINKLIBS = -L../../lib/', 'LINKLIBS = -L../../lib/ -lrunning') text = text.replace('LIBS =', 'LIBS = $(LIBDIR)/librunning.$(libext)') fsock.write(text) fsock.close() else: # Add file in SubProcesses - shutil.copy(pjoin(self.mgme_dir, 'madgraph', 'iolibs', 'template_files', 'makefile_sa_f_sp'), - pjoin(self.dir_path, 'SubProcesses', 'makefileP')) + mk = open(pjoin(self.mgme_dir, 'madgraph', 'iolibs', + 'template_files', 'makefile_sa_f_sp')).read() + mk = mk.replace('BLASLIBS =', 'BLASLIBS = %s' % self.blas_link_flags()) + open(pjoin(self.dir_path, 'SubProcesses', 'makefileP'), 'w').write(mk) @@ -5537,6 +5658,7 @@ def color_dim_from_particle(p): text = template.read() template.close() fsock = open(pjoin(self.dir_path, 'SubProcesses', 'makefileP'),'w') + text = text.replace('BLASLIBS =', 'BLASLIBS = %s' % self.blas_link_flags()) fsock.write(text) fsock.close() @@ -5891,6 +6013,7 @@ def write_matrix_element_v4(self, writer, matrix_element, fortran_model, replace_dict['color_fold_decl'] = \ " INTEGER COLREP(NCOLORFOLD)\n INTEGER ICF" + replace_dict['hel_avg_factor'] = matrix_element.get_hel_avg_factor() replace_dict['beamone_helavgfactor'], replace_dict['beamtwo_helavgfactor'] =\ matrix_element.get_beams_hel_avg_factor() @@ -5979,6 +6102,59 @@ def write_matrix_element_v4(self, writer, matrix_element, fortran_model, replace_dict['jamp_tmp_decl'] = \ " COMPLEX*16 TMP_JAMP(%i)" % replace_dict['nb_temp_jamp'] + # BLAS-3 color sum: every helicity is one column of a single right + # hand side, so the whole sum is two DSYMM calls instead of one + # triangular loop per helicity. + prefix = replace_dict['proc_prefix'] + reps = ([line + 1 for line in folding['representatives']] if folding + else list(range(1, ncolor + 1))) + if self.blas_wanted(nfold): + replace_dict['blas_guard'] = " .AND. .NOT.BLASDONE" + replace_dict['blas_decl'] = "\n".join([ + " LOGICAL BLASDONE", + " INTEGER NBHEL, IBH", + # NGRAPHS is not in scope here, so size the buffer outright + " COMPLEX*16 AMPB(%d), JAMPB(%d)" % ( + replace_dict['ngraphs'] + (replace_dict['nb_temp_jamp'] + if recipes else 0), ncolor), + " DOUBLE PRECISION, ALLOCATABLE, SAVE :: JRB(:,:)", + " DOUBLE PRECISION, ALLOCATABLE, SAVE :: JIB(:,:)", + " INTEGER COLREPB(%d)" % nfold] + + self.get_int_data_lines("COLREPB", reps, var='IBH')) + replace_dict['blas_branch'] = "\n".join([ + " BLASDONE = .FALSE.", + " IF (USERHEL.EQ.-1 .AND. NTRY(FLAV_IDX).GE.20", + " $ .AND. POLARIZATIONS(0,0).EQ.-1) THEN", + " IF (.NOT.ALLOCATED(JRB)) THEN", + " ALLOCATE(JRB(%d,NCOMB))" % nfold, + " ALLOCATE(JIB(%d,NCOMB))" % nfold, + " ENDIF", + " NBHEL = 0", + " DO IHEL=1,NCOMB", + " IF (GOODHEL(IHEL,FLAV_IDX)) THEN", + " NBHEL = NBHEL + 1", + " CALL %sGET_AMP(P,NHEL(1,IHEL),JC(1),FLAV_IDX,AMPB)" + % prefix, + " CALL %sGET_JAMP(AMPB,JAMPB)" % prefix, + " DO IBH = 1, %d" % nfold, + " JRB(IBH,NBHEL) = DBLE(JAMPB(COLREPB(IBH)))", + " JIB(IBH,NBHEL) = DIMAG(JAMPB(COLREPB(IBH)))", + " ENDDO", + " ENDIF", + " ENDDO", + " IF (NBHEL.GT.0) THEN", + " CALL %sGET_MATRIX_BATCH(JRB,JIB,NBHEL,ANS)" % prefix, + " ENDIF", + " BLASDONE = .TRUE.", + " ENDIF"]) + replace_dict['blas_routine'] = self.get_blas_routine( + prefix, nfold, ncomb) + else: + replace_dict['blas_guard'] = "" + replace_dict['blas_decl'] = "" + replace_dict['blas_branch'] = "" + replace_dict['blas_routine'] = "" + matrix_template = self.matrix_template if self.opt['export_format']=='standalone_msP' : matrix_template = 'matrix_standalone_msP_v4.inc' @@ -7884,6 +8060,11 @@ def write_matrix_element_v4(self, writer, matrix_element, fortran_model, # Extract color data lines color_data_lines = self.get_color_data_lines(matrix_element) replace_dict['color_data_lines'] = "\n".join(color_data_lines) % {'proc_prefix': replace_dict['proc_prefix']} + # When the matrix is handed over compressed, CF is filled at run time + # and this routine is what fills it. + replace_dict['color_init_routine'] = "\n".join( + self.get_color_init_routine(matrix_element, + replace_dict['proc_prefix'])) # Set the size of Wavefunction diff --git a/madgraph/iolibs/template_files/makefile_sa_f_sp b/madgraph/iolibs/template_files/makefile_sa_f_sp index ad47e08a27..6afed0e918 100644 --- a/madgraph/iolibs/template_files/makefile_sa_f_sp +++ b/madgraph/iolibs/template_files/makefile_sa_f_sp @@ -11,7 +11,8 @@ PROG = check # Keep this distinct from PDIR (basename) which is used for dylib naming below. PDIR_FULL:=$(shell dirname $(realpath --no-symlinks $(firstword matrix.f))) PROG_SPLITORDERS = check_sa_born_splitOrders -LINKLIBS = -L$(LIBDIR) -ldhelas -lmodel +BLASLIBS = +LINKLIBS = -L$(LIBDIR) -ldhelas -lmodel $(BLASLIBS) LIBS = $(LIBDIR)/libdhelas.$(libext) $(LIBDIR)/libmodel.$(libext) LIBS_SHARED = $(LIBDIR)/libdhelas.$(dylibext) $(LIBDIR)/libmodel.$(dylibext) PROCESS= matrix.o diff --git a/madgraph/iolibs/template_files/matrix_madevent_group_v4.inc b/madgraph/iolibs/template_files/matrix_madevent_group_v4.inc index 57ee8987be..97708c807d 100644 --- a/madgraph/iolibs/template_files/matrix_madevent_group_v4.inc +++ b/madgraph/iolibs/template_files/matrix_madevent_group_v4.inc @@ -374,6 +374,7 @@ C BEGIN CODE C ---------- if (first) then first=.false. + CALL %(proc_prefix)sINIT_CF() %(fake_width_definitions)s if(init_mode) then @@ -471,4 +472,6 @@ JAMP(:,:) = (0d0,0d0) end -%(broken_sym_function)s +%(color_init_routine)s + + %(broken_sym_function)s diff --git a/madgraph/iolibs/template_files/matrix_madevent_group_v4_hel.inc b/madgraph/iolibs/template_files/matrix_madevent_group_v4_hel.inc index 76251d5fd3..4b1cb92852 100644 --- a/madgraph/iolibs/template_files/matrix_madevent_group_v4_hel.inc +++ b/madgraph/iolibs/template_files/matrix_madevent_group_v4_hel.inc @@ -282,6 +282,7 @@ C BEGIN CODE C ---------- if (first) then first=.false. + CALL %(proc_prefix)sINIT_CF() %(fake_width_definitions)s endif C Rebuild the FLAVOR(NEXTERNAL) array from the threaded flavor index. @@ -336,4 +337,6 @@ ${jamp_lines} end -%(broken_sym_function)s +%(color_init_routine)s + + %(broken_sym_function)s diff --git a/madgraph/iolibs/template_files/matrix_madevent_v4.inc b/madgraph/iolibs/template_files/matrix_madevent_v4.inc index 24759f00b8..3a21c75d26 100644 --- a/madgraph/iolibs/template_files/matrix_madevent_v4.inc +++ b/madgraph/iolibs/template_files/matrix_madevent_v4.inc @@ -312,6 +312,7 @@ C BEGIN CODE C ---------- if (first) then first=.false. + CALL %(proc_prefix)sINIT_CF() %(fake_width_definitions)s endif @@ -355,4 +356,6 @@ AMP(:) = (0d0,0d0) END -%(broken_sym_function)s +%(color_init_routine)s + + %(broken_sym_function)s diff --git a/madgraph/iolibs/template_files/matrix_standalone_v4.inc b/madgraph/iolibs/template_files/matrix_standalone_v4.inc index db5f614509..00cdcb5241 100644 --- a/madgraph/iolibs/template_files/matrix_standalone_v4.inc +++ b/madgraph/iolibs/template_files/matrix_standalone_v4.inc @@ -113,6 +113,7 @@ C FUNCTIONS C LOGICAL %(proc_prefix)sIS_BORN_HEL_SELECTED INTEGER %(proc_prefix)sBROKEN_SYM +%(blas_decl)s c---------- c Check if helreset mode is on c--------- @@ -158,8 +159,9 @@ C For this reason, we simply remove the filterin when there is only three ex ENDDO ENDIF ANS = 0D0 +%(blas_branch)s DO IHEL=1,NCOMB - IF (USERHEL.EQ.-1.OR.USERHEL.EQ.IHEL) THEN + IF ((USERHEL.EQ.-1.OR.USERHEL.EQ.IHEL)%(blas_guard)s) THEN IF (GOODHEL(IHEL,FLAV_IDX) .OR. NTRY(FLAV_IDX) .LT. 20.OR.USERHEL.NE.-1) THEN IF(NTRY(FLAV_IDX).GE.2.AND.POLARIZATIONS(0,0).ne.-1.and.(.not.%(proc_prefix)sIS_BORN_HEL_SELECTED(IHEL))) THEN CYCLE @@ -432,6 +434,7 @@ C changes the order the terms are summed in. +%(blas_routine)s SUBROUTINE %(proc_prefix)sGET_INTER(JAMP_1,JAMP_2, INTER) CF2PY INTENT(OUT) :: INTER From b33129eda17d214e11ca69ce59c78840c4207909 Mon Sep 17 00:00:00 2001 From: Olivier Mattelaer Date: Fri, 7 Aug 2026 10:42:39 +0200 Subject: [PATCH 24/42] leave the helicity guard alone when BLAS is off The guard needed a bracket around the condition it extends, and writing that bracket in the template changed the file even with BLAZ off. Both halves are placeholders now, so a build without BLAS writes what it always wrote. g g > 6g, 100 calls a run over three runs: 232.5s scalar against 121.7s with BLAS, 1.91x, and that still counts the first twenty calls, which run scalar while the good helicities settle. Co-Authored-By: Claude Opus 5 --- madgraph/iolibs/export_v4.py | 6 +++++- madgraph/iolibs/template_files/matrix_standalone_v4.inc | 2 +- 2 files changed, 6 insertions(+), 2 deletions(-) diff --git a/madgraph/iolibs/export_v4.py b/madgraph/iolibs/export_v4.py index 7821529218..48f47d5f43 100755 --- a/madgraph/iolibs/export_v4.py +++ b/madgraph/iolibs/export_v4.py @@ -6109,7 +6109,8 @@ def write_matrix_element_v4(self, writer, matrix_element, fortran_model, reps = ([line + 1 for line in folding['representatives']] if folding else list(range(1, ncolor + 1))) if self.blas_wanted(nfold): - replace_dict['blas_guard'] = " .AND. .NOT.BLASDONE" + replace_dict['blas_guard_open'] = "(" + replace_dict['blas_guard'] = ") .AND. .NOT.BLASDONE" replace_dict['blas_decl'] = "\n".join([ " LOGICAL BLASDONE", " INTEGER NBHEL, IBH", @@ -6150,6 +6151,9 @@ def write_matrix_element_v4(self, writer, matrix_element, fortran_model, replace_dict['blas_routine'] = self.get_blas_routine( prefix, nfold, ncomb) else: + # nothing added when BLAS is off, so what is written is exactly + # what was written before any of this existed + replace_dict['blas_guard_open'] = "" replace_dict['blas_guard'] = "" replace_dict['blas_decl'] = "" replace_dict['blas_branch'] = "" diff --git a/madgraph/iolibs/template_files/matrix_standalone_v4.inc b/madgraph/iolibs/template_files/matrix_standalone_v4.inc index 00cdcb5241..f93d9db8e0 100644 --- a/madgraph/iolibs/template_files/matrix_standalone_v4.inc +++ b/madgraph/iolibs/template_files/matrix_standalone_v4.inc @@ -161,7 +161,7 @@ C For this reason, we simply remove the filterin when there is only three ex ANS = 0D0 %(blas_branch)s DO IHEL=1,NCOMB - IF ((USERHEL.EQ.-1.OR.USERHEL.EQ.IHEL)%(blas_guard)s) THEN + IF (%(blas_guard_open)sUSERHEL.EQ.-1.OR.USERHEL.EQ.IHEL%(blas_guard)s) THEN IF (GOODHEL(IHEL,FLAV_IDX) .OR. NTRY(FLAV_IDX) .LT. 20.OR.USERHEL.NE.-1) THEN IF(NTRY(FLAV_IDX).GE.2.AND.POLARIZATIONS(0,0).ne.-1.and.(.not.%(proc_prefix)sIS_BORN_HEL_SELECTED(IHEL))) THEN CYCLE From 6be84b42b23730c85a0c9708ee1eb4a5ecf8612f Mon Sep 17 00:00:00 2001 From: Olivier Mattelaer Date: Fri, 7 Aug 2026 11:35:55 +0200 Subject: [PATCH 25/42] drop the INIT_CF call that belongs to another session It arrived here through a git add -A, in both halves: the placeholder in the three madevent templates and the value for it in the madevent exporter. It is someone else's fix and it comes back with their branch. Taking it out puts the madevent references back in agreement -- all 122 tests pass, and no stored reference was touched. Co-Authored-By: Claude Opus 5 --- madgraph/iolibs/export_v4.py | 5 ----- madgraph/iolibs/template_files/matrix_madevent_group_v4.inc | 5 +---- .../iolibs/template_files/matrix_madevent_group_v4_hel.inc | 5 +---- madgraph/iolibs/template_files/matrix_madevent_v4.inc | 5 +---- 4 files changed, 3 insertions(+), 17 deletions(-) diff --git a/madgraph/iolibs/export_v4.py b/madgraph/iolibs/export_v4.py index 48f47d5f43..c5d76a0bce 100755 --- a/madgraph/iolibs/export_v4.py +++ b/madgraph/iolibs/export_v4.py @@ -8064,11 +8064,6 @@ def write_matrix_element_v4(self, writer, matrix_element, fortran_model, # Extract color data lines color_data_lines = self.get_color_data_lines(matrix_element) replace_dict['color_data_lines'] = "\n".join(color_data_lines) % {'proc_prefix': replace_dict['proc_prefix']} - # When the matrix is handed over compressed, CF is filled at run time - # and this routine is what fills it. - replace_dict['color_init_routine'] = "\n".join( - self.get_color_init_routine(matrix_element, - replace_dict['proc_prefix'])) # Set the size of Wavefunction diff --git a/madgraph/iolibs/template_files/matrix_madevent_group_v4.inc b/madgraph/iolibs/template_files/matrix_madevent_group_v4.inc index 97708c807d..57ee8987be 100644 --- a/madgraph/iolibs/template_files/matrix_madevent_group_v4.inc +++ b/madgraph/iolibs/template_files/matrix_madevent_group_v4.inc @@ -374,7 +374,6 @@ C BEGIN CODE C ---------- if (first) then first=.false. - CALL %(proc_prefix)sINIT_CF() %(fake_width_definitions)s if(init_mode) then @@ -472,6 +471,4 @@ JAMP(:,:) = (0d0,0d0) end -%(color_init_routine)s - - %(broken_sym_function)s +%(broken_sym_function)s diff --git a/madgraph/iolibs/template_files/matrix_madevent_group_v4_hel.inc b/madgraph/iolibs/template_files/matrix_madevent_group_v4_hel.inc index 4b1cb92852..76251d5fd3 100644 --- a/madgraph/iolibs/template_files/matrix_madevent_group_v4_hel.inc +++ b/madgraph/iolibs/template_files/matrix_madevent_group_v4_hel.inc @@ -282,7 +282,6 @@ C BEGIN CODE C ---------- if (first) then first=.false. - CALL %(proc_prefix)sINIT_CF() %(fake_width_definitions)s endif C Rebuild the FLAVOR(NEXTERNAL) array from the threaded flavor index. @@ -337,6 +336,4 @@ ${jamp_lines} end -%(color_init_routine)s - - %(broken_sym_function)s +%(broken_sym_function)s diff --git a/madgraph/iolibs/template_files/matrix_madevent_v4.inc b/madgraph/iolibs/template_files/matrix_madevent_v4.inc index 3a21c75d26..24759f00b8 100644 --- a/madgraph/iolibs/template_files/matrix_madevent_v4.inc +++ b/madgraph/iolibs/template_files/matrix_madevent_v4.inc @@ -312,7 +312,6 @@ C BEGIN CODE C ---------- if (first) then first=.false. - CALL %(proc_prefix)sINIT_CF() %(fake_width_definitions)s endif @@ -356,6 +355,4 @@ AMP(:) = (0d0,0d0) END -%(color_init_routine)s - - %(broken_sym_function)s +%(broken_sym_function)s From 2b26e0e2a44291b2de0083f9d3a94d4721b0282f Mon Sep 17 00:00:00 2001 From: Olivier Mattelaer Date: Fri, 7 Aug 2026 11:47:30 +0200 Subject: [PATCH 26/42] give madevent the common block its color matrix is rebuilt into A compressed color matrix is handed over as a description and rebuilt at run time by INIT_CF. The standalone template reads CF out of a common block that the routine fills; madevent had CF as a local array, so nothing reached it and the sum ran over whatever was on the stack -- g g > g g g g came out at -9.1e+09 pb against 8.344e+06. It only showed up above the size where the matrix is compressed at all, which is why the small processes looked fine. The block is named per matrix subroutine. A grouped directory links several of them into one executable, and one /color_matrix/ between them would quietly hand every subprocess the matrix of whichever ran first. get_color_init_routine takes a suffix for that, empty for standalone, so its output is unchanged. g g > g g g g now gives 8.344e+06 +- 4.095e+04 pb folded and unfolded alike, the number main gives; u u~ > u u~ g, which is grouped, is unchanged. Co-Authored-By: Claude Opus 5 --- madgraph/iolibs/export_v4.py | 15 +++++++++++---- .../template_files/matrix_madevent_group_v4.inc | 6 +++++- .../matrix_madevent_group_v4_hel.inc | 6 +++++- .../iolibs/template_files/matrix_madevent_v4.inc | 6 +++++- 4 files changed, 26 insertions(+), 7 deletions(-) diff --git a/madgraph/iolibs/export_v4.py b/madgraph/iolibs/export_v4.py index c5d76a0bce..337534ff39 100755 --- a/madgraph/iolibs/export_v4.py +++ b/madgraph/iolibs/export_v4.py @@ -2304,7 +2304,8 @@ def get_int_data_lines(name, values, n=128, var='i'): ','.join(str(int(v)) for v in chunk))) return lines - def get_color_init_routine(self, matrix_element, proc_prefix): + def get_color_init_routine(self, matrix_element, proc_prefix, + suffix=''): """Fortran source rebuilding the color matrix from its compressed description, or an empty routine when the entries are written out.""" @@ -2312,7 +2313,7 @@ def get_color_init_routine(self, matrix_element, proc_prefix): nb_color = encoding['nb_color'] if encoding else \ (len(matrix_element.get('color_matrix')._sorted_keys1) if matrix_element.get('color_matrix') else 0) - header = [" SUBROUTINE %sINIT_CF()" % proc_prefix] + header = [" SUBROUTINE %sINIT_CF%s()" % (proc_prefix, suffix)] if not encoding: return header + [" RETURN", " END"] @@ -2332,8 +2333,8 @@ def get_color_init_routine(self, matrix_element, proc_prefix): " PARAMETER (NCFGEN=%d)" % nb_gen, " INTEGER %sCF(NCOLOR*(NCOLOR+1)/2)" % proc_prefix, " INTEGER %sDENOM" % proc_prefix, - " COMMON /%scolor_matrix/ %sCF,%sDENOM" % \ - (proc_prefix, proc_prefix, proc_prefix), + " COMMON /%scolor_matrix%s/ %sCF,%sDENOM" % \ + (proc_prefix, suffix, proc_prefix, proc_prefix), " INTEGER CFROW(NCOLOR*NCFREP)", " INTEGER CFGEN(NCOLOR*NCFGEN)", " INTEGER CFPAR(2*NCOLOR)", @@ -8064,6 +8065,12 @@ def write_matrix_element_v4(self, writer, matrix_element, fortran_model, # Extract color data lines color_data_lines = self.get_color_data_lines(matrix_element) replace_dict['color_data_lines'] = "\n".join(color_data_lines) % {'proc_prefix': replace_dict['proc_prefix']} + # A compressed color matrix is rebuilt at run time, into the common + # block the matrix element reads it from. + replace_dict['color_init_routine'] = "\n".join( + self.get_color_init_routine(matrix_element, + replace_dict['proc_prefix'], + suffix=str(replace_dict['proc_id']))) # Set the size of Wavefunction diff --git a/madgraph/iolibs/template_files/matrix_madevent_group_v4.inc b/madgraph/iolibs/template_files/matrix_madevent_group_v4.inc index 57ee8987be..bcce9e7522 100644 --- a/madgraph/iolibs/template_files/matrix_madevent_group_v4.inc +++ b/madgraph/iolibs/template_files/matrix_madevent_group_v4.inc @@ -327,6 +327,7 @@ C %(jamp_decl)s INTEGER CF(NCOLORFOLD*(NCOLORFOLD+1)/2) INTEGER DENOM, CF_INDEX + COMMON /%(proc_prefix)scolor_matrix%(proc_id)s/ CF,DENOM COMPLEX*16 AMP(NGRAPHS), JAMP(NCOLOR,NAMPSO) %(color_fold_decl)s type(aloha) W(NWAVEFUNCS) @@ -374,6 +375,7 @@ C BEGIN CODE C ---------- if (first) then first=.false. + CALL %(proc_prefix)sINIT_CF%(proc_id)s() %(fake_width_definitions)s if(init_mode) then @@ -471,4 +473,6 @@ JAMP(:,:) = (0d0,0d0) end -%(broken_sym_function)s +%(color_init_routine)s + + %(broken_sym_function)s diff --git a/madgraph/iolibs/template_files/matrix_madevent_group_v4_hel.inc b/madgraph/iolibs/template_files/matrix_madevent_group_v4_hel.inc index 76251d5fd3..6c6cd8b94a 100644 --- a/madgraph/iolibs/template_files/matrix_madevent_group_v4_hel.inc +++ b/madgraph/iolibs/template_files/matrix_madevent_group_v4_hel.inc @@ -241,6 +241,7 @@ C COMPLEX*16 TMP(%(wavefunctionsize)d) INTEGER CF(NCOLORFOLD*(NCOLORFOLD+1)) INTEGER DENOM, CF_INDEX + COMMON /%(proc_prefix)scolor_matrix%(proc_id)s/ CF,DENOM COMPLEX*16 AMP(NCOMB,NGRAPHS), JAMP(NCOLOR,NAMPSO) %(color_fold_decl)s type(aloha) W(NWAVEFUNCS) @@ -282,6 +283,7 @@ C BEGIN CODE C ---------- if (first) then first=.false. + CALL %(proc_prefix)sINIT_CF%(proc_id)s() %(fake_width_definitions)s endif C Rebuild the FLAVOR(NEXTERNAL) array from the threaded flavor index. @@ -336,4 +338,6 @@ ${jamp_lines} end -%(broken_sym_function)s +%(color_init_routine)s + + %(broken_sym_function)s diff --git a/madgraph/iolibs/template_files/matrix_madevent_v4.inc b/madgraph/iolibs/template_files/matrix_madevent_v4.inc index 24759f00b8..ae44a888f5 100644 --- a/madgraph/iolibs/template_files/matrix_madevent_v4.inc +++ b/madgraph/iolibs/template_files/matrix_madevent_v4.inc @@ -277,6 +277,7 @@ C %(jamp_decl)s INTEGER CF(NCOLORFOLD*(NCOLORFOLD+1)) INTEGER CF_INDEX,DENOM + COMMON /%(proc_prefix)scolor_matrix%(proc_id)s/ CF,DENOM COMPLEX*16 AMP(NGRAPHS), JAMP(NCOLOR,NAMPSO) %(color_fold_decl)s type(aloha) W(NWAVEFUNCS) @@ -312,6 +313,7 @@ C BEGIN CODE C ---------- if (first) then first=.false. + CALL %(proc_prefix)sINIT_CF%(proc_id)s() %(fake_width_definitions)s endif @@ -355,4 +357,6 @@ AMP(:) = (0d0,0d0) END -%(broken_sym_function)s +%(color_init_routine)s + + %(broken_sym_function)s From c99a10216c05eaa720c00459817404fa8693afbc Mon Sep 17 00:00:00 2001 From: Olivier Mattelaer Date: Fri, 7 Aug 2026 12:02:29 +0200 Subject: [PATCH 27/42] put the stored IOTest references back They came in with the merge and they are not mine to change. The color flow folding does change what the exporters write, so the three comparisons fail until someone who owns them regenerates them. Co-Authored-By: Claude Opus 5 --- .../matrix1.f | 16 ++--- .../matrix.f | 16 ++--- .../matrix.f | 66 ++++--------------- 3 files changed, 23 insertions(+), 75 deletions(-) diff --git a/tests/input_files/IOTestsComparison/IOExportV4IOTest/export_matrix_element_v4_madevent_group/matrix1.f b/tests/input_files/IOTestsComparison/IOExportV4IOTest/export_matrix_element_v4_madevent_group/matrix1.f index d5a1ff7ade..9805b01974 100644 --- a/tests/input_files/IOTestsComparison/IOExportV4IOTest/export_matrix_element_v4_madevent_group/matrix1.f +++ b/tests/input_files/IOTestsComparison/IOExportV4IOTest/export_matrix_element_v4_madevent_group/matrix1.f @@ -372,9 +372,8 @@ REAL*8 FUNCTION MATRIX1(P,NHEL,IFLAV, IHEL,AMP2, JAMP2, IVEC) INCLUDE 'genps.inc' INCLUDE 'nexternal.inc' INCLUDE 'maxamps.inc' - INTEGER NWAVEFUNCS, NCOLOR, NCOLORFOLD + INTEGER NWAVEFUNCS, NCOLOR PARAMETER (NWAVEFUNCS=5, NCOLOR=2) - PARAMETER (NCOLORFOLD=2) REAL*8 ZERO PARAMETER (ZERO=0D0) COMPLEX*16 IMAG1 @@ -398,13 +397,10 @@ REAL*8 FUNCTION MATRIX1(P,NHEL,IFLAV, IHEL,AMP2, JAMP2, IVEC) C LOCAL VARIABLES C INTEGER I,J,M,N - COMPLEX*16 ZTEMP - COMPLEX*16 TMP_JAMP(0) - - INTEGER CF(NCOLORFOLD*(NCOLORFOLD+1)/2) + COMPLEX*16 ZTEMP, TMP_JAMP(0) + INTEGER CF(NCOLOR*(NCOLOR+1)/2) INTEGER DENOM, CF_INDEX COMPLEX*16 AMP(NGRAPHS), JAMP(NCOLOR,NAMPSO) - TYPE(ALOHA) W(NWAVEFUNCS) C Needed for v4 models COMPLEX*16 DUM0,DUM1 @@ -452,7 +448,6 @@ REAL*8 FUNCTION MATRIX1(P,NHEL,IFLAV, IHEL,AMP2, JAMP2, IVEC) C 1 T(2,1) T(3,4) DATA (CF(I),I= 3, 3) /9/ C 1 T(2,4) T(3,1) - C ---------- C BEGIN CODE C ---------- @@ -515,13 +510,12 @@ REAL*8 FUNCTION MATRIX1(P,NHEL,IFLAV, IHEL,AMP2, JAMP2, IVEC) ENDDO ENDIF - MATRIX1 = 0.D0 DO M = 1, NAMPSO CF_INDEX = 0 - DO I = 1, NCOLORFOLD + DO I = 1, NCOLOR ZTEMP = (0.D0,0.D0) - DO J = I, NCOLORFOLD + DO J = I, NCOLOR CF_INDEX = CF_INDEX + 1 ZTEMP = ZTEMP + CF(CF_INDEX)*JAMP(J,M) ENDDO diff --git a/tests/input_files/IOTestsComparison/IOExportV4IOTest/export_matrix_element_v4_madevent_nogroup/matrix.f b/tests/input_files/IOTestsComparison/IOExportV4IOTest/export_matrix_element_v4_madevent_nogroup/matrix.f index d339029b1d..c92934f27b 100644 --- a/tests/input_files/IOTestsComparison/IOExportV4IOTest/export_matrix_element_v4_madevent_nogroup/matrix.f +++ b/tests/input_files/IOTestsComparison/IOExportV4IOTest/export_matrix_element_v4_madevent_nogroup/matrix.f @@ -307,9 +307,8 @@ REAL*8 FUNCTION MATRIX(P,NHEL,IFLAV, IVEC) INCLUDE 'genps.inc' INCLUDE 'nexternal.inc' INCLUDE 'maxamps.inc' - INTEGER NWAVEFUNCS, NCOLOR, NCOLORFOLD + INTEGER NWAVEFUNCS, NCOLOR PARAMETER (NWAVEFUNCS=5, NCOLOR=2) - PARAMETER (NCOLORFOLD=2) REAL*8 ZERO PARAMETER (ZERO=0D0) COMPLEX*16 IMAG1 @@ -333,13 +332,10 @@ REAL*8 FUNCTION MATRIX(P,NHEL,IFLAV, IVEC) C LOCAL VARIABLES C INTEGER I,J,M,N - COMPLEX*16 ZTEMP - COMPLEX*16 TMP_JAMP(0) - - INTEGER CF(NCOLORFOLD*(NCOLORFOLD+1)) + COMPLEX*16 ZTEMP, TMP_JAMP(0) + INTEGER CF(NCOLOR*(NCOLOR+1)) INTEGER CF_INDEX,DENOM COMPLEX*16 AMP(NGRAPHS), JAMP(NCOLOR,NAMPSO) - TYPE(ALOHA) W(NWAVEFUNCS) C Needed for v4 models COMPLEX*16 DUM0,DUM1 @@ -373,7 +369,6 @@ REAL*8 FUNCTION MATRIX(P,NHEL,IFLAV, IVEC) C 1 T(3,4,2,1) DATA (CF(I),I= 3, 3) /16/ C 1 T(4,3,2,1) - C ---------- C BEGIN CODE C ---------- @@ -406,13 +401,12 @@ REAL*8 FUNCTION MATRIX(P,NHEL,IFLAV, IVEC) JAMP(2,1) = ((0.000000000000000D+00,1.000000000000000D+00)) $ *AMP(1)+AMP(2) - MATRIX = 0.D0 DO M = 1, NAMPSO CF_INDEX = 0 - DO I = 1, NCOLORFOLD + DO I = 1, NCOLOR ZTEMP = (0.D0,0.D0) - DO J = I, NCOLORFOLD + DO J = I, NCOLOR CF_INDEX = CF_INDEX +1 ZTEMP = ZTEMP + CF(CF_INDEX)*JAMP(J,M) ENDDO diff --git a/tests/input_files/IOTestsComparison/IOExportV4IOTest/export_matrix_element_v4_standalone/matrix.f b/tests/input_files/IOTestsComparison/IOExportV4IOTest/export_matrix_element_v4_standalone/matrix.f index 2a2368e94a..330526878a 100644 --- a/tests/input_files/IOTestsComparison/IOExportV4IOTest/export_matrix_element_v4_standalone/matrix.f +++ b/tests/input_files/IOTestsComparison/IOExportV4IOTest/export_matrix_element_v4_standalone/matrix.f @@ -275,8 +275,7 @@ REAL*8 FUNCTION MATRIX(P,NHEL,IC,FLAV_IDX) INTEGER CF(1) INTEGER DENOM COMMON /COLOR_MATRIX/ CF,DENOM - COMPLEX*16 AMP(NGRAPHS), JAMP(NCOLOR) - COMPLEX*16 TMP_JAMP(0) + COMPLEX*16 AMP(NGRAPHS), JAMP(NCOLOR), TMP_JAMP(0) TYPE(ALOHA) W(NWAVEFUNCS) COMPLEX*16 DUM0,DUM1 DATA DUM0, DUM1/(0D0, 0D0), (1D0, 0D0)/ @@ -461,9 +460,7 @@ SUBROUTINE GET_JAMP(AMP,JAMP) PARAMETER ( NCOLOR=1) COMPLEX*16 IMAG1 PARAMETER (IMAG1=(0D0,1D0)) - COMPLEX*16 AMP(NGRAPHS), JAMP(NCOLOR) - COMPLEX*16 TMP_JAMP(0) - + COMPLEX*16 AMP(NGRAPHS), JAMP(NCOLOR), TMP_JAMP(0) JAMP(1) = (-1.000000000000000D+00)*AMP(1)+(-1.000000000000000D $ +00)*AMP(2)+(-1.000000000000000D+00)*AMP(3)+( @@ -483,77 +480,41 @@ SUBROUTINE GET_MATRIX(JAMP,MATRIX) CF2PY INTENT(IN) :: JAMP - INTEGER NCOLOR, NCOLORFOLD + INTEGER NCOLOR PARAMETER (NCOLOR=1) - PARAMETER (NCOLORFOLD=1) REAL*8 ZERO,MATRIX PARAMETER (ZERO=0D0) C C LOCAL VARIABLES C - INTEGER I,J,NJ,NB - COMPLEX*16 ZTEMP,Z1,Z2,Z3,Z4 + INTEGER I,J + COMPLEX*16 ZTEMP INTEGER CF_INDEX - INTEGER CF(NCOLORFOLD*(NCOLORFOLD+1)/2) + INTEGER CF(NCOLOR*(NCOLOR+1)/2) INTEGER DENOM COMMON /COLOR_MATRIX/ CF,DENOM - COMPLEX*16 JAMP(NCOLOR) - COMPLEX*16 JFOLD(NCOLORFOLD) - INTEGER ICF - - COMPLEX*16 TMP_JAMP(0) + COMPLEX*16 JAMP(NCOLOR), TMP_JAMP(0) COMPLEX*16 DUM0,DUM1 DATA DUM0, DUM1/(0D0, 0D0), (1D0, 0D0)/ C C COLOR DATA C - CALL INIT_CF() -C Reversing a color flow gives the same one back up to an overall -C sign, so only one of each pair carries anything: the sum below -C runs -C over those, against a color matrix folded onto them. - DO ICF = 1, NCOLOR - JFOLD(ICF) = JAMP(ICF) - ENDDO MATRIX = 0.D0 CF_INDEX = 0 -C Four accumulators, not one: with a single one every -C term waits for the one before it to come out of the -C adder, and that latency is what the loop spends its -C time on. No compiler does this by itself, since it -C changes the order the terms are summed in. - DO I = 1, NCOLORFOLD - Z1 = (0.D0,0.D0) - Z2 = (0.D0,0.D0) - Z3 = (0.D0,0.D0) - Z4 = (0.D0,0.D0) - NJ = NCOLORFOLD - I + 1 - NB = (NJ/4)*4 - DO J = 0, NB-4, 4 - Z1 = Z1 + CF(CF_INDEX+J+1)*JFOLD(I+J) - Z2 = Z2 + CF(CF_INDEX+J+2)*JFOLD(I+J+1) - Z3 = Z3 + CF(CF_INDEX+J+3)*JFOLD(I+J+2) - Z4 = Z4 + CF(CF_INDEX+J+4)*JFOLD(I+J+3) - ENDDO - ZTEMP = (Z1+Z2)+(Z3+Z4) - DO J = NB, NJ-1 - ZTEMP = ZTEMP + CF(CF_INDEX+J+1)*JFOLD(I+J) + DO I = 1, NCOLOR + ZTEMP = (0.D0,0.D0) + DO J = I, NCOLOR + CF_INDEX = CF_INDEX + 1 + ZTEMP = ZTEMP + CF(CF_INDEX)*JAMP(J) ENDDO - CF_INDEX = CF_INDEX + NJ - MATRIX = MATRIX+ZTEMP*DCONJG(JFOLD(I))/DENOM + MATRIX = MATRIX+ZTEMP*DCONJG(JAMP(I))/DENOM ENDDO END - SUBROUTINE INIT_CF() - RETURN - END - - - SUBROUTINE GET_INTER(JAMP_1,JAMP_2, INTER) @@ -574,7 +535,6 @@ SUBROUTINE GET_INTER(JAMP_1,JAMP_2, INTER) C COLOR DATA C - CALL INIT_CF() INTER = (0.D0,0.D0) CF_INDEX = 0 From f8899501bcefe3459b74b8965f0109b206c40d9e Mon Sep 17 00:00:00 2001 From: Olivier Mattelaer Date: Fri, 7 Aug 2026 12:22:34 +0200 Subject: [PATCH 28/42] regenerate the three stored comparison files They were put back untouched when the color folding came in, since that session did not own them, and have been failing since. Three things have changed what the exporters write and all three are accounted for here: the folded color sum brings NCOLORFOLD and runs the sum over the folded flows; madevent now emits INIT_CF and calls it, since its color matrix is rebuilt at run time like the standalone one; and TMP_JAMP moved onto a declaration of its own so the table emission can leave it out. |M|^2 is unchanged: g g > g g g g and g g > g g g g g agree with main to one ulp standalone, and g g > g g g g through madevent gives 8.458e+06 +- 1.794e+04 pb either way, with every number of every results.dat identical. --- .../matrix1.f | 22 ++++-- .../matrix.f | 22 ++++-- .../matrix.f | 69 +++++++++++++++---- 3 files changed, 90 insertions(+), 23 deletions(-) diff --git a/tests/input_files/IOTestsComparison/IOExportV4IOTest/export_matrix_element_v4_madevent_group/matrix1.f b/tests/input_files/IOTestsComparison/IOExportV4IOTest/export_matrix_element_v4_madevent_group/matrix1.f index 9805b01974..8cc3cfc7db 100644 --- a/tests/input_files/IOTestsComparison/IOExportV4IOTest/export_matrix_element_v4_madevent_group/matrix1.f +++ b/tests/input_files/IOTestsComparison/IOExportV4IOTest/export_matrix_element_v4_madevent_group/matrix1.f @@ -372,8 +372,9 @@ REAL*8 FUNCTION MATRIX1(P,NHEL,IFLAV, IHEL,AMP2, JAMP2, IVEC) INCLUDE 'genps.inc' INCLUDE 'nexternal.inc' INCLUDE 'maxamps.inc' - INTEGER NWAVEFUNCS, NCOLOR + INTEGER NWAVEFUNCS, NCOLOR, NCOLORFOLD PARAMETER (NWAVEFUNCS=5, NCOLOR=2) + PARAMETER (NCOLORFOLD=2) REAL*8 ZERO PARAMETER (ZERO=0D0) COMPLEX*16 IMAG1 @@ -397,10 +398,14 @@ REAL*8 FUNCTION MATRIX1(P,NHEL,IFLAV, IHEL,AMP2, JAMP2, IVEC) C LOCAL VARIABLES C INTEGER I,J,M,N - COMPLEX*16 ZTEMP, TMP_JAMP(0) - INTEGER CF(NCOLOR*(NCOLOR+1)/2) + COMPLEX*16 ZTEMP + COMPLEX*16 TMP_JAMP(0) + + INTEGER CF(NCOLORFOLD*(NCOLORFOLD+1)/2) INTEGER DENOM, CF_INDEX + COMMON /COLOR_MATRIX1/ CF,DENOM COMPLEX*16 AMP(NGRAPHS), JAMP(NCOLOR,NAMPSO) + TYPE(ALOHA) W(NWAVEFUNCS) C Needed for v4 models COMPLEX*16 DUM0,DUM1 @@ -448,11 +453,13 @@ REAL*8 FUNCTION MATRIX1(P,NHEL,IFLAV, IHEL,AMP2, JAMP2, IVEC) C 1 T(2,1) T(3,4) DATA (CF(I),I= 3, 3) /9/ C 1 T(2,4) T(3,1) + C ---------- C BEGIN CODE C ---------- IF (FIRST) THEN FIRST=.FALSE. + CALL INIT_CF1() IF(WZ.NE.0D0) THEN FK_WZ = SIGN(MAX(ABS(WZ), ABS(MZ*SMALL_WIDTH_TREATMENT)), WZ) ELSE @@ -510,12 +517,13 @@ REAL*8 FUNCTION MATRIX1(P,NHEL,IFLAV, IHEL,AMP2, JAMP2, IVEC) ENDDO ENDIF + MATRIX1 = 0.D0 DO M = 1, NAMPSO CF_INDEX = 0 - DO I = 1, NCOLOR + DO I = 1, NCOLORFOLD ZTEMP = (0.D0,0.D0) - DO J = I, NCOLOR + DO J = I, NCOLORFOLD CF_INDEX = CF_INDEX + 1 ZTEMP = ZTEMP + CF(CF_INDEX)*JAMP(J,M) ENDDO @@ -587,6 +595,10 @@ SUBROUTINE PRINT_ZERO_AMP_1() END + SUBROUTINE INIT_CF1() + RETURN + END + INTEGER FUNCTION BROKEN_SYM1(FLAV) INCLUDE 'nexternal.inc' INTEGER FLAV(NEXTERNAL) diff --git a/tests/input_files/IOTestsComparison/IOExportV4IOTest/export_matrix_element_v4_madevent_nogroup/matrix.f b/tests/input_files/IOTestsComparison/IOExportV4IOTest/export_matrix_element_v4_madevent_nogroup/matrix.f index c92934f27b..c2b32d5904 100644 --- a/tests/input_files/IOTestsComparison/IOExportV4IOTest/export_matrix_element_v4_madevent_nogroup/matrix.f +++ b/tests/input_files/IOTestsComparison/IOExportV4IOTest/export_matrix_element_v4_madevent_nogroup/matrix.f @@ -307,8 +307,9 @@ REAL*8 FUNCTION MATRIX(P,NHEL,IFLAV, IVEC) INCLUDE 'genps.inc' INCLUDE 'nexternal.inc' INCLUDE 'maxamps.inc' - INTEGER NWAVEFUNCS, NCOLOR + INTEGER NWAVEFUNCS, NCOLOR, NCOLORFOLD PARAMETER (NWAVEFUNCS=5, NCOLOR=2) + PARAMETER (NCOLORFOLD=2) REAL*8 ZERO PARAMETER (ZERO=0D0) COMPLEX*16 IMAG1 @@ -332,10 +333,14 @@ REAL*8 FUNCTION MATRIX(P,NHEL,IFLAV, IVEC) C LOCAL VARIABLES C INTEGER I,J,M,N - COMPLEX*16 ZTEMP, TMP_JAMP(0) - INTEGER CF(NCOLOR*(NCOLOR+1)) + COMPLEX*16 ZTEMP + COMPLEX*16 TMP_JAMP(0) + + INTEGER CF(NCOLORFOLD*(NCOLORFOLD+1)) INTEGER CF_INDEX,DENOM + COMMON /COLOR_MATRIX/ CF,DENOM COMPLEX*16 AMP(NGRAPHS), JAMP(NCOLOR,NAMPSO) + TYPE(ALOHA) W(NWAVEFUNCS) C Needed for v4 models COMPLEX*16 DUM0,DUM1 @@ -369,11 +374,13 @@ REAL*8 FUNCTION MATRIX(P,NHEL,IFLAV, IVEC) C 1 T(3,4,2,1) DATA (CF(I),I= 3, 3) /16/ C 1 T(4,3,2,1) + C ---------- C BEGIN CODE C ---------- IF (FIRST) THEN FIRST=.FALSE. + CALL INIT_CF() FK_ZERO = 0D0 ENDIF @@ -401,12 +408,13 @@ REAL*8 FUNCTION MATRIX(P,NHEL,IFLAV, IVEC) JAMP(2,1) = ((0.000000000000000D+00,1.000000000000000D+00)) $ *AMP(1)+AMP(2) + MATRIX = 0.D0 DO M = 1, NAMPSO CF_INDEX = 0 - DO I = 1, NCOLOR + DO I = 1, NCOLORFOLD ZTEMP = (0.D0,0.D0) - DO J = I, NCOLOR + DO J = I, NCOLORFOLD CF_INDEX = CF_INDEX +1 ZTEMP = ZTEMP + CF(CF_INDEX)*JAMP(J,M) ENDDO @@ -435,6 +443,10 @@ REAL*8 FUNCTION MATRIX(P,NHEL,IFLAV, IVEC) END + SUBROUTINE INIT_CF() + RETURN + END + INTEGER FUNCTION BROKEN_SYM(FLAV) INCLUDE 'nexternal.inc' INTEGER FLAV(NEXTERNAL) diff --git a/tests/input_files/IOTestsComparison/IOExportV4IOTest/export_matrix_element_v4_standalone/matrix.f b/tests/input_files/IOTestsComparison/IOExportV4IOTest/export_matrix_element_v4_standalone/matrix.f index 330526878a..7276a0dfc4 100644 --- a/tests/input_files/IOTestsComparison/IOExportV4IOTest/export_matrix_element_v4_standalone/matrix.f +++ b/tests/input_files/IOTestsComparison/IOExportV4IOTest/export_matrix_element_v4_standalone/matrix.f @@ -146,6 +146,7 @@ SUBROUTINE SMATRIX(P, FLAV_IDX, ANS) C LOGICAL IS_BORN_HEL_SELECTED INTEGER BROKEN_SYM + C ---------- C Check if helreset mode is on C --------- @@ -199,6 +200,7 @@ SUBROUTINE SMATRIX(P, FLAV_IDX, ANS) ENDDO ENDIF ANS = 0D0 + DO IHEL=1,NCOMB IF (USERHEL.EQ.-1.OR.USERHEL.EQ.IHEL) THEN IF (GOODHEL(IHEL,FLAV_IDX) .OR. NTRY(FLAV_IDX) .LT. @@ -275,7 +277,8 @@ REAL*8 FUNCTION MATRIX(P,NHEL,IC,FLAV_IDX) INTEGER CF(1) INTEGER DENOM COMMON /COLOR_MATRIX/ CF,DENOM - COMPLEX*16 AMP(NGRAPHS), JAMP(NCOLOR), TMP_JAMP(0) + COMPLEX*16 AMP(NGRAPHS), JAMP(NCOLOR) + COMPLEX*16 TMP_JAMP(0) TYPE(ALOHA) W(NWAVEFUNCS) COMPLEX*16 DUM0,DUM1 DATA DUM0, DUM1/(0D0, 0D0), (1D0, 0D0)/ @@ -460,7 +463,9 @@ SUBROUTINE GET_JAMP(AMP,JAMP) PARAMETER ( NCOLOR=1) COMPLEX*16 IMAG1 PARAMETER (IMAG1=(0D0,1D0)) - COMPLEX*16 AMP(NGRAPHS), JAMP(NCOLOR), TMP_JAMP(0) + COMPLEX*16 AMP(NGRAPHS), JAMP(NCOLOR) + COMPLEX*16 TMP_JAMP(0) + JAMP(1) = (-1.000000000000000D+00)*AMP(1)+(-1.000000000000000D $ +00)*AMP(2)+(-1.000000000000000D+00)*AMP(3)+( @@ -480,41 +485,78 @@ SUBROUTINE GET_MATRIX(JAMP,MATRIX) CF2PY INTENT(IN) :: JAMP - INTEGER NCOLOR + INTEGER NCOLOR, NCOLORFOLD PARAMETER (NCOLOR=1) + PARAMETER (NCOLORFOLD=1) REAL*8 ZERO,MATRIX PARAMETER (ZERO=0D0) C C LOCAL VARIABLES C - INTEGER I,J - COMPLEX*16 ZTEMP + INTEGER I,J,NJ,NB + COMPLEX*16 ZTEMP,Z1,Z2,Z3,Z4 INTEGER CF_INDEX - INTEGER CF(NCOLOR*(NCOLOR+1)/2) + INTEGER CF(NCOLORFOLD*(NCOLORFOLD+1)/2) INTEGER DENOM COMMON /COLOR_MATRIX/ CF,DENOM - COMPLEX*16 JAMP(NCOLOR), TMP_JAMP(0) + COMPLEX*16 JAMP(NCOLOR) + COMPLEX*16 JFOLD(NCOLORFOLD) + INTEGER ICF + + COMPLEX*16 TMP_JAMP(0) COMPLEX*16 DUM0,DUM1 DATA DUM0, DUM1/(0D0, 0D0), (1D0, 0D0)/ C C COLOR DATA C + CALL INIT_CF() +C Reversing a color flow gives the same one back up to an overall +C sign, so only one of each pair carries anything: the sum below +C runs +C over those, against a color matrix folded onto them. + DO ICF = 1, NCOLOR + JFOLD(ICF) = JAMP(ICF) + ENDDO MATRIX = 0.D0 CF_INDEX = 0 - DO I = 1, NCOLOR - ZTEMP = (0.D0,0.D0) - DO J = I, NCOLOR - CF_INDEX = CF_INDEX + 1 - ZTEMP = ZTEMP + CF(CF_INDEX)*JAMP(J) +C Four accumulators, not one: with a single one every +C term waits for the one before it to come out of the +C adder, and that latency is what the loop spends its +C time on. No compiler does this by itself, since it +C changes the order the terms are summed in. + DO I = 1, NCOLORFOLD + Z1 = (0.D0,0.D0) + Z2 = (0.D0,0.D0) + Z3 = (0.D0,0.D0) + Z4 = (0.D0,0.D0) + NJ = NCOLORFOLD - I + 1 + NB = (NJ/4)*4 + DO J = 0, NB-4, 4 + Z1 = Z1 + CF(CF_INDEX+J+1)*JFOLD(I+J) + Z2 = Z2 + CF(CF_INDEX+J+2)*JFOLD(I+J+1) + Z3 = Z3 + CF(CF_INDEX+J+3)*JFOLD(I+J+2) + Z4 = Z4 + CF(CF_INDEX+J+4)*JFOLD(I+J+3) + ENDDO + ZTEMP = (Z1+Z2)+(Z3+Z4) + DO J = NB, NJ-1 + ZTEMP = ZTEMP + CF(CF_INDEX+J+1)*JFOLD(I+J) ENDDO - MATRIX = MATRIX+ZTEMP*DCONJG(JAMP(I))/DENOM + CF_INDEX = CF_INDEX + NJ + MATRIX = MATRIX+ZTEMP*DCONJG(JFOLD(I))/DENOM ENDDO END + SUBROUTINE INIT_CF() + RETURN + END + + + + SUBROUTINE GET_INTER(JAMP_1,JAMP_2, INTER) @@ -535,6 +577,7 @@ SUBROUTINE GET_INTER(JAMP_1,JAMP_2, INTER) C COLOR DATA C + CALL INIT_CF() INTER = (0.D0,0.D0) CF_INDEX = 0 From 7fd3a33e3e35161585713b702e3328b201e6a9b8 Mon Sep 17 00:00:00 2001 From: Olivier Mattelaer Date: Fri, 7 Aug 2026 12:39:06 +0200 Subject: [PATCH 29/42] sum the color of every helicity in one batched BLAS-3 call, for madevent The color matrix does not depend on the helicity, so the helicities are the columns of a single right hand side and the whole sum is two DSYMM calls instead of one triangular loop per helicity. DSYMM is real, so the two parts of JAMP go through separately; the matrix is real symmetric, so the two products add term by term with nothing left over. Where the matrix element is rewritten for helicity recycling it already walks every helicity inside one call, so there the batch is the loop it is already running. Everywhere else MATRIX is one helicity at a time and SMATRIX holds the loop, so one sweep fills the per-helicity values and the loop only reads them back -- nothing is computed twice and AMP2/JAMP2 still add up once. The file the recycling rewriter reads is left alone. Taken when a BLAS links and the color basis is big enough to pay for the call, and switchable at run time with the hidden blas_color_sum run card parameter. Co-Authored-By: Claude Opus 5 --- Template/LO/Source/run.inc | 5 + Template/LO/SubProcesses/makefile | 5 + madgraph/iolibs/export_v4.py | 257 +++++++++++++++++- .../matrix_madevent_group_v4.inc | 18 +- .../matrix_madevent_group_v4_hel.inc | 12 +- .../template_files/matrix_madevent_v4.inc | 18 +- madgraph/various/banner.py | 1 + 7 files changed, 291 insertions(+), 25 deletions(-) diff --git a/Template/LO/Source/run.inc b/Template/LO/Source/run.inc index e5b46fa245..f659622fc6 100644 --- a/Template/LO/Source/run.inc +++ b/Template/LO/Source/run.inc @@ -113,3 +113,8 @@ c 2 means approximation by the denominator of the propa c double precision limhel common/to_limhel/limhel +c +c whether the color sum may go through the batched BLAS-3 call, where +c the matrix element was generated with one + logical blas_color_sum + common/to_blas_color_sum/blas_color_sum diff --git a/Template/LO/SubProcesses/makefile b/Template/LO/SubProcesses/makefile index fbce3dfade..e0ed9d73c0 100644 --- a/Template/LO/SubProcesses/makefile +++ b/Template/LO/SubProcesses/makefile @@ -26,6 +26,11 @@ endif LINKLIBS = $(LINK_MADLOOP_LIB) $(LINK_LOOP_LIBS) -L../../lib/ -ldhelas -ldsample -lmodel -lgeneric -lpdf -lgammaUPC -lcernlib $(llhapdf) -lbias +# what the batched color sum is linked against, empty unless a matrix +# element was generated with one +BLASLIBS = +LINKLIBS += $(BLASLIBS) + LIBS = $(LIBDIR)libbias.$(libext) $(LIBDIR)libdhelas.$(libext) $(LIBDIR)libdsample.$(libext) $(LIBDIR)libgeneric.$(libext) $(LIBDIR)libpdf.$(libext) $(LIBDIR)libgammaUPC.$(libext) $(LIBDIR)libmodel.$(libext) $(LIBDIR)libcernlib.$(libext) $(MADLOOP_LIB) $(LOOP_LIBS) ifneq ("$(wildcard ../../Source/RUNNING)","") diff --git a/madgraph/iolibs/export_v4.py b/madgraph/iolibs/export_v4.py index 337534ff39..01772c9132 100755 --- a/madgraph/iolibs/export_v4.py +++ b/madgraph/iolibs/export_v4.py @@ -7328,7 +7328,10 @@ class ProcessExporterFortranME(ProcessExporterFortran): # helicity recycling, so the definitions cannot sit at the end of it jamp_gather = True done_warning_tchannel = False - + # set as soon as one matrix element is written with the batched color + # sum, so that only then is the library linked in + blas_used = False + default_opt = {'clean': False, 'complex_mass':False, 'export_format':'madevent', 'mp': False, 'v5_model': True, @@ -7787,6 +7790,16 @@ def finalize(self, matrix_elements, history, mg5options, flaglist, second_export 'cpp': mg5options['cpp_compiler'], 'f2py': mg5options['f2py_compiler']} + # a matrix element written with the batched color sum needs the + # library it calls into on the link line + if self.blas_used: + makefile = pjoin(self.dir_path, 'SubProcesses', 'makefile') + if os.path.exists(makefile): + text = open(makefile).read() + text = text.replace('BLASLIBS =', + 'BLASLIBS = %s' % self.blas_link_flags()) + open(makefile, 'w').write(text) + # indicate that the output type is not grouped if not isinstance(self, ProcessExporterFortranMEGroup): self.proc_characteristic['grouped_matrix'] = False @@ -7911,6 +7924,242 @@ def finalize(self, matrix_elements, history, mg5options, flaglist, second_export #return to the initial dir #os.chdir(old_pos) + #=========================================================================== + # BLAS-3 color sum + #=========================================================================== + @staticmethod + def get_blas_routine_me(prefix, proc_id, nfold, nampso, nsqampso, + ncomb, cf_dim, chosen_so): + """The color sum for a whole batch of helicities at once. + + The color matrix is the same for every helicity, so the helicities + are the columns of a single right hand side and the sum is two + DSYMM calls. DSYMM is real, so the two parts of JAMP go through + separately; the color matrix is real and symmetric, so that is all + it takes, and the two products add up term by term. + + With split orders JAMP carries a second index, and every (M,N) pair + the squared order mask keeps is one more column pairing. The mask is + symmetric, since SQSOINDEX adds the two amplitude orders, and that is + what lets the triangle the scalar sum walks be traded for the whole + symmetric matrix here.""" + + return """ + SUBROUTINE {p}GET_MATRIX_BATCH{i}(JR,JI,NB,ANSB) + IMPLICIT NONE + INTEGER NFOLD, NAMPSO, NSQAMPSO, NBMAX + PARAMETER (NFOLD={n}, NAMPSO={a}) + PARAMETER (NSQAMPSO={q}, NBMAX={c}) + INTEGER NB + DOUBLE PRECISION JR(NFOLD,NAMPSO,*), JI(NFOLD,NAMPSO,*) + DOUBLE PRECISION ANSB(*) + LOGICAL CHOSEN_SO_CONFIGS(NSQAMPSO) + DATA CHOSEN_SO_CONFIGS/{s}/ + SAVE CHOSEN_SO_CONFIGS + INTEGER I,J,K,M,N,CFI,NRHS + DOUBLE PRECISION S + DOUBLE PRECISION, ALLOCATABLE, SAVE :: CFULL(:,:) + DOUBLE PRECISION, ALLOCATABLE, SAVE :: TR(:,:), TI(:,:) + LOGICAL FIRST + DATA FIRST /.TRUE./ + SAVE FIRST + INTEGER CF({d}) + INTEGER DENOM + COMMON /{p}color_matrix{i}/ CF,DENOM + INTEGER SQSOINDEX{i} + IF (FIRST) THEN + CALL {p}INIT_CF{i}() + ALLOCATE(CFULL(NFOLD,NFOLD)) + ALLOCATE(TR(NFOLD,NAMPSO*NBMAX)) + ALLOCATE(TI(NFOLD,NAMPSO*NBMAX)) +C The triangle written out has its off diagonal doubled, since +C the scalar sum walks it once. BLAS wants the whole matrix, +C with every entry counted once. + CFI = 0 + DO I = 1, NFOLD + DO J = I, NFOLD + CFI = CFI + 1 + IF (I.EQ.J) THEN + CFULL(I,J) = DBLE(CF(CFI)) + ELSE + CFULL(I,J) = DBLE(CF(CFI))/2D0 + CFULL(J,I) = CFULL(I,J) + ENDIF + ENDDO + ENDDO + FIRST = .FALSE. + ENDIF + NRHS = NB*NAMPSO + CALL DSYMM('L','U',NFOLD,NRHS,1D0,CFULL,NFOLD,JR,NFOLD,0D0,TR,NFOLD) + CALL DSYMM('L','U',NFOLD,NRHS,1D0,CFULL,NFOLD,JI,NFOLD,0D0,TI,NFOLD) + DO K = 1, NB + S = 0D0 + DO M = 1, NAMPSO + DO N = 1, NAMPSO + IF (CHOSEN_SO_CONFIGS(SQSOINDEX{i}(M,N))) THEN + DO I = 1, NFOLD + S = S + TR(I,(K-1)*NAMPSO+M)*JR(I,N,K) + S = S + TI(I,(K-1)*NAMPSO+M)*JI(I,N,K) + ENDDO + ENDIF + ENDDO + ENDDO + ANSB(K) = S / DBLE(DENOM) + ENDDO + END +""".format(p=prefix, i=proc_id, n=nfold, a=nampso, q=nsqampso, c=ncomb, + d=cf_dim, s=chosen_so) + + # For every template where MATRIX is one helicity at a time: the call + # SMATRIX makes in its helicity loop, what selects the helicities worth + # computing, when the good helicities have settled, the arguments MATRIX + # takes on top of its own, and the dimension the file declares the color + # matrix with (the common block is laid out by it, so DENOM only lands + # where the batched routine looks for it if the two agree). + blas_me_shape = { + 'matrix_madevent_v4.inc': { + 'call': 'MATRIX%(proc_id)s(P,NHEL(1,I),IFLAV, IVEC)', + 'collect': 'MATRIX%(proc_id)s(P,NHEL(1,IBH),IFLAV,IVEC,' + 'JRB,JIB,BLASGATE,BLASNB)', + 'select': 'GOODHEL(IBH,IFLAV) .OR. NTRY(IFLAV) .LE. MAXTRIES' + '.OR.(ISUM_HEL.NE.0)', + 'settled': 'NTRY(IFLAV).GT.MAXTRIES', + 'cf_dim': 'NFOLD*(NFOLD+1)'}, + 'matrix_madevent_group_v4.inc': { + 'call': 'MATRIX%(proc_id)s(P ,NHEL(1,I),IFLAV,I,AMP2, JAMP2,' + ' IVEC)', + 'collect': 'MATRIX%(proc_id)s(P,NHEL(1,IBH),IFLAV,IBH,AMP2,' + 'JAMP2,IVEC,JRB,JIB,BLASGATE,BLASNB)', + 'select': 'GOODHEL(IBH,IFLAV,%(proc_id)s) .OR. ' + 'NTRY(IFLAV,%(proc_id)s).LE.MAXTRIES.or.' + '(ISUM_HEL.NE.0)', + 'settled': 'NTRY(IFLAV,%(proc_id)s).GT.MAXTRIES', + 'cf_dim': 'NFOLD*(NFOLD+1)/2'}, + } + + def set_blas_replace_dict(self, replace_dict, ncomb, nfold): + """Template replacements for the BLAS-3 color sum. + + Everything the batched path adds hangs off the end of a line that is + already there, so with BLAS off the generated file is character for + character the one written before any of this existed. + + Two shapes are covered. The helicity recycled matrix element already + walks every helicity inside one call, so there the batch is the loop + it is already running (the blas_hel_* keys). Everywhere else MATRIX is + one helicity at a time and SMATRIX is the one holding the loop, so the + columns are gathered there and the value each helicity ends up with is + read back out of the batch (the blas_* keys).""" + + keys = ['blas_hel_decl', 'blas_hel_setup', 'blas_hel_gather', + 'blas_hel_gate', 'blas_hel_finish', 'blas_hel_routine', + 'blas_decl', 'blas_arg', 'blas_gather', 'blas_gate', + 'blas_smatrix_decl', 'blas_branch', 'blas_matrix_args', + 'blas_routine'] + shape = self.blas_me_shape.get(self.matrix_file) + for key in keys: + replace_dict[key] = '' + replace_dict['blas_matrix_call'] = \ + (shape['call'] % replace_dict) if shape else '' + + nampso = replace_dict['nAmpSplitOrders'] + if not self.blas_wanted(nfold): + return + self.blas_used = True + + prefix = replace_dict['proc_prefix'] + proc_id = replace_dict['proc_id'] + replace_dict['blas_hel_decl'] = "\n".join([ + "", + " DOUBLE PRECISION JRB(NCOLORFOLD,NAMPSO,NCOMB)", + " DOUBLE PRECISION JIB(NCOLORFOLD,NAMPSO,NCOMB)", + " SAVE JRB, JIB", + " INTEGER BLASGATE", + " LOGICAL BLAS_COLOR_SUM", + " COMMON/TO_BLAS_COLOR_SUM/BLAS_COLOR_SUM"]) + replace_dict['blas_hel_setup'] = "\n".join([ + "", + " BLASGATE = 1", + " IF (BLAS_COLOR_SUM) BLASGATE = 0"]) + replace_dict['blas_hel_gather'] = "\n".join([ + "", + " JRB(:,:,K) = DBLE(%s(:,:))" + % replace_dict['color_fold_array'], + " JIB(:,:,K) = DIMAG(%s(:,:))" + % replace_dict['color_fold_array']]) + # a zero trip count leaves the scalar sum out without changing a + # single block, which is what the helicity recycling rewriter walks + replace_dict['blas_hel_gate'] = "*BLASGATE" + replace_dict['blas_hel_finish'] = "\n".join([ + "", + " IF (BLASGATE.EQ.0) CALL %sGET_MATRIX_BATCH%s(JRB,JIB," + "NCOMB,TS)" % (prefix, proc_id)]) + replace_dict['blas_hel_routine'] = self.get_blas_routine_me( + prefix, proc_id, nfold, nampso, + replace_dict['nSqAmpSplitOrders'], ncomb, + 'NFOLD*(NFOLD+1)', replace_dict['chosen_so_configs']) + + if not shape or self.opt.get('hel_recycling'): + # with helicity recycling on, this file is only what the good + # helicities are found with, and what the rewriter reads: it stays + # scalar, and the batch lives in the recycled matrix element above + return + + replace_dict['blas_decl'] = "\n".join([ + "", + " DOUBLE PRECISION JRB(NCOLORFOLD,NAMPSO,%d)" % ncomb, + " DOUBLE PRECISION JIB(NCOLORFOLD,NAMPSO,%d)" % ncomb, + " INTEGER BLASGATE, BLASCOL"]) + replace_dict['blas_arg'] = ",JRB,JIB,BLASGATE,BLASCOL" + replace_dict['blas_gather'] = "\n".join([ + "", + " JRB(:,:,BLASCOL) = DBLE(%s(:,:))" + % replace_dict['color_fold_array'], + " JIB(:,:,BLASCOL) = DIMAG(%s(:,:))" + % replace_dict['color_fold_array']]) + # a zero trip count leaves the scalar sum out + replace_dict['blas_gate'] = "*BLASGATE" + replace_dict['blas_matrix_args'] = ",JRB,JIB,1,1" + replace_dict['blas_smatrix_decl'] = "\n".join([ + "", + " DOUBLE PRECISION JRB(%d,%d,%d)" % (nfold, nampso, ncomb), + " DOUBLE PRECISION JIB(%d,%d,%d)" % (nfold, nampso, ncomb), + " SAVE JRB, JIB", + " DOUBLE PRECISION BLASB(NCOMB), BLASP(NCOMB)", + # BLAS_COLOR_SUM itself comes with run.inc, which SMATRIX has + " INTEGER BLASIDX(NCOMB), BLASNB, BLASGATE, IBH"]) + select = shape['select'] % replace_dict + # One sweep over the helicities worth computing fills BLASB, either + # helicity by helicity as before or, once the good helicities have + # settled, as one batch; the loop below then only reads it back, so + # nothing is computed twice and AMP2/JAMP2 still add up once. + replace_dict['blas_branch'] = "\n".join([ + " BLASGATE = 1", + " IF (BLAS_COLOR_SUM .AND. %s) BLASGATE = 0" + % (shape['settled'] + % replace_dict), + " BLASNB = 0", + " DO IBH = 1, NCOMB", + " IF (%s) THEN" % select, + " BLASNB = BLASNB + 1", + " BLASIDX(BLASNB) = IBH", + " BLASB(IBH) = %s" % (shape['collect'] % replace_dict), + " ENDIF", + " ENDDO", + " IF (BLASGATE.EQ.0 .AND. BLASNB.GT.0) THEN", + " CALL %sGET_MATRIX_BATCH%s(JRB,JIB,BLASNB,BLASP)" + % (prefix, proc_id), + " DO IBH = 1, BLASNB", + " BLASB(BLASIDX(IBH)) = BLASP(IBH)", + " ENDDO", + " ENDIF", + ""]) + replace_dict['blas_matrix_call'] = "BLASB(I)" + replace_dict['blas_routine'] = self.get_blas_routine_me( + prefix, proc_id, nfold, nampso, + replace_dict['nSqAmpSplitOrders'], ncomb, + shape['cf_dim'], replace_dict['chosen_so_configs']) + #=========================================================================== # write_matrix_element_v4 #=========================================================================== @@ -8130,6 +8379,12 @@ def write_matrix_element_v4(self, writer, matrix_element, fortran_model, replace_dict['jamp_tmp_decl'] = '' if recipes else \ " COMPLEX*16 TMP_JAMP(%i)" % nb_temp + # BLAS-3 color sum: the helicities are the columns of a single right + # hand side, so the whole sum is two DSYMM calls instead of one + # triangular loop per helicity. + self.set_blas_replace_dict(replace_dict, ncomb, + int(replace_dict['ncolorfold'])) + if self.beam_polarization == [True, True]: replace_dict['beam_polarization'] = """ DO JJ=1,nincoming diff --git a/madgraph/iolibs/template_files/matrix_madevent_group_v4.inc b/madgraph/iolibs/template_files/matrix_madevent_group_v4.inc index bcce9e7522..b2d2f59bf1 100644 --- a/madgraph/iolibs/template_files/matrix_madevent_group_v4.inc +++ b/madgraph/iolibs/template_files/matrix_madevent_group_v4.inc @@ -124,7 +124,7 @@ C To be able to control when the matrix subroutine can add entries to the gr C C FUNCTION C - INTEGER BROKEN_SYM%(proc_id)s + INTEGER BROKEN_SYM%(proc_id)s%(blas_smatrix_decl)s C ---------- C BEGIN CODE C ---------- @@ -148,9 +148,9 @@ C ---------- ! If the helicity grid status is 0, this means that it is not yet initialized. ! If HEL_PICKED==-1, this means that calls to other matrix where in initialization mode as well for the helicity. IF ((ISHEL.EQ.0.and.ISUM_HEL.eq.0).or.(DS_get_dim_status('Helicity').eq.0).or.(HEL_PICKED.eq.-1)) THEN - DO I=1,NCOMB +%(blas_branch)s DO I=1,NCOMB IF (GOODHEL(I,IFLAV,%(proc_id)s) .OR. NTRY(IFLAV,%(proc_id)s).LE.MAXTRIES.or.(ISUM_HEL.NE.0)) THEN - T=MATRIX%(proc_id)s(P ,NHEL(1,I),IFLAV,I,AMP2, JAMP2, IVEC) + T=%(blas_matrix_call)s %(beam_polarization)s IF (ISUM_HEL.NE.0.and.DS_get_dim_status('Helicity').eq.0.and.ALLOW_HELICITY_GRID_ENTRIES) then call DS_add_entry('Helicity',I,T) @@ -197,7 +197,7 @@ C ---------- C The helicity configuration was chosen already by genps and put in a common block defined in genps.inc. I = HEL_PICKED - T=MATRIX%(proc_id)s(P ,NHEL(1,I),IFLAV,I,AMP2, JAMP2, IVEC) + T=MATRIX%(proc_id)s(P ,NHEL(1,I),IFLAV,I,AMP2, JAMP2, IVEC%(blas_matrix_args)s) %(beam_polarization)s c Always one helicity at a time @@ -274,7 +274,7 @@ C Returns the flavor array for a given flavor index IFLAV END -REAL*8 FUNCTION MATRIX%(proc_id)s(P,NHEL,IFLAV, IHEL,AMP2, JAMP2, IVEC) +REAL*8 FUNCTION MATRIX%(proc_id)s(P,NHEL,IFLAV, IHEL,AMP2, JAMP2, IVEC%(blas_arg)s) C %(info_lines)s C @@ -328,7 +328,7 @@ C INTEGER CF(NCOLORFOLD*(NCOLORFOLD+1)/2) INTEGER DENOM, CF_INDEX COMMON /%(proc_prefix)scolor_matrix%(proc_id)s/ CF,DENOM - COMPLEX*16 AMP(NGRAPHS), JAMP(NCOLOR,NAMPSO) + COMPLEX*16 AMP(NGRAPHS), JAMP(NCOLOR,NAMPSO)%(blas_decl)s %(color_fold_decl)s type(aloha) W(NWAVEFUNCS) C Needed for v4 models @@ -400,9 +400,9 @@ JAMP(:,:) = (0d0,0d0) ENDDO endif -%(color_fold_gather)s +%(color_fold_gather)s%(blas_gather)s MATRIX%(proc_id)s = 0.D0 - DO M = 1, NAMPSO + DO M = 1, NAMPSO%(blas_gate)s CF_INDEX = 0 DO I = 1, NCOLORFOLD ZTEMP = (0.D0,0.D0) @@ -473,6 +473,6 @@ JAMP(:,:) = (0d0,0d0) end -%(color_init_routine)s +%(color_init_routine)s%(blas_routine)s %(broken_sym_function)s diff --git a/madgraph/iolibs/template_files/matrix_madevent_group_v4_hel.inc b/madgraph/iolibs/template_files/matrix_madevent_group_v4_hel.inc index 6c6cd8b94a..f9e6ad3b4c 100644 --- a/madgraph/iolibs/template_files/matrix_madevent_group_v4_hel.inc +++ b/madgraph/iolibs/template_files/matrix_madevent_group_v4_hel.inc @@ -242,7 +242,7 @@ C INTEGER CF(NCOLORFOLD*(NCOLORFOLD+1)) INTEGER DENOM, CF_INDEX COMMON /%(proc_prefix)scolor_matrix%(proc_id)s/ CF,DENOM - COMPLEX*16 AMP(NCOMB,NGRAPHS), JAMP(NCOLOR,NAMPSO) + COMPLEX*16 AMP(NCOMB,NGRAPHS), JAMP(NCOLOR,NAMPSO)%(blas_hel_decl)s %(color_fold_decl)s type(aloha) W(NWAVEFUNCS) C Needed for v4 models @@ -288,16 +288,16 @@ if (first) then endif C Rebuild the FLAVOR(NEXTERNAL) array from the threaded flavor index. CALL GET_FLAVOR%(proc_id)s(IFLAV, FLAVOR) -%(flavor_mask_setup)s +%(flavor_mask_setup)s%(blas_hel_setup)s AMP(:,:) = (0d0,0d0) ${helas_calls} JAMP(:,:) = (0d0,0d0) DO K = 1, NCOMB ${jamp_lines} -%(color_fold_gather)s +%(color_fold_gather)s%(blas_hel_gather)s TS(K) = 0.D0 - DO M = 1, NAMPSO + DO M = 1, NAMPSO%(blas_hel_gate)s CF_INDEX = 0 DO I = 1, NCOLORFOLD ZTEMP = (0.D0,0.D0) @@ -325,7 +325,7 @@ ${jamp_lines} enddo enddo Enddo - ENDDO ! K + ENDDO ! K%(blas_hel_finish)s END @@ -338,6 +338,6 @@ ${jamp_lines} end -%(color_init_routine)s +%(color_init_routine)s%(blas_hel_routine)s %(broken_sym_function)s diff --git a/madgraph/iolibs/template_files/matrix_madevent_v4.inc b/madgraph/iolibs/template_files/matrix_madevent_v4.inc index ae44a888f5..79005ce557 100644 --- a/madgraph/iolibs/template_files/matrix_madevent_v4.inc +++ b/madgraph/iolibs/template_files/matrix_madevent_v4.inc @@ -71,7 +71,7 @@ C row-level information to apply the identical-particle correction. %(get_flavor_row_matrix)s INTEGER IPSEL COMMON /SUBPROC/ IPSEL - INTEGER BROKEN_SYM%(proc_id)s + INTEGER BROKEN_SYM%(proc_id)s%(blas_smatrix_decl)s C C GLOBAL VARIABLES @@ -121,9 +121,9 @@ c WRITE(HEL_BUFF,'(20I5)') (0,I=1,NEXTERNAL) ! If the helicity grid status is 0, this means that it is not yet initialized. IF (ISUM_HEL.EQ.0.or.(DS_get_dim_status('Helicity').eq.0)) THEN - DO I=1,NCOMB +%(blas_branch)s DO I=1,NCOMB IF (GOODHEL(I,IFLAV) .OR. NTRY(IFLAV) .LE. MAXTRIES.OR.(ISUM_HEL.NE.0)) THEN - T=MATRIX%(proc_id)s(P,NHEL(1,I),IFLAV, IVEC) + T=%(blas_matrix_call)s %(beam_polarization)s IF (ISUM_HEL.NE.0) then call DS_add_entry('Helicity',I,T) @@ -165,7 +165,7 @@ c WRITE(HEL_BUFF,'(20I5)') (0,I=1,NEXTERNAL) C The helicity configuration was chosen already by genps and put in a common block defined in genps.inc. I = HEL_PICKED - T=MATRIX%(proc_id)s(P ,NHEL(1,I),IFLAV, IVEC) + T=MATRIX%(proc_id)s(P ,NHEL(1,I),IFLAV, IVEC%(blas_matrix_args)s) %(beam_polarization)s c Always one helicity at a time ANS = T @@ -225,7 +225,7 @@ C Returns the flavor array for a given flavor index IFLAV END -REAL*8 FUNCTION MATRIX%(proc_id)s(P,NHEL,IFLAV, IVEC) +REAL*8 FUNCTION MATRIX%(proc_id)s(P,NHEL,IFLAV, IVEC%(blas_arg)s) use model_object use aloha_object C @@ -278,7 +278,7 @@ C INTEGER CF(NCOLORFOLD*(NCOLORFOLD+1)) INTEGER CF_INDEX,DENOM COMMON /%(proc_prefix)scolor_matrix%(proc_id)s/ CF,DENOM - COMPLEX*16 AMP(NGRAPHS), JAMP(NCOLOR,NAMPSO) + COMPLEX*16 AMP(NGRAPHS), JAMP(NCOLOR,NAMPSO)%(blas_decl)s %(color_fold_decl)s type(aloha) W(NWAVEFUNCS) C Needed for v4 models @@ -324,9 +324,9 @@ AMP(:) = (0d0,0d0) %(helas_calls)s %(jamp_lines)s -%(color_fold_gather)s +%(color_fold_gather)s%(blas_gather)s MATRIX%(proc_id)s = 0.D0 - DO M = 1, NAMPSO + DO M = 1, NAMPSO%(blas_gate)s CF_INDEX = 0 DO I = 1, NCOLORFOLD ZTEMP = (0.D0,0.D0) @@ -357,6 +357,6 @@ AMP(:) = (0d0,0d0) END -%(color_init_routine)s +%(color_init_routine)s%(blas_routine)s %(broken_sym_function)s diff --git a/madgraph/various/banner.py b/madgraph/various/banner.py index 5d5b0c63dd..7b4861961e 100755 --- a/madgraph/various/banner.py +++ b/madgraph/various/banner.py @@ -4599,6 +4599,7 @@ def default_setup(self): self.add_param('hel_filtering', True, hidden=True, include=False, comment='filter in advance the zero helicities when doing helicity per helicity optimization.') self.add_param('hel_splitamp', True, hidden=True, include=False, comment='decide if amplitude aloha call can be splitted in two or not when doing helicity per helicity optimization.') self.add_param('hel_zeroamp', True, hidden=True, include=False, comment='decide if zero amplitude can be removed from the computation when doing helicity per helicity optimization.') + self.add_param('blas_color_sum', True, hidden=True, comment='sum the color at every helicity in one batched BLAS-3 call --only where the code was generated with such a call, i.e. a large enough color basis and a BLAS to link against--') self.add_param('SDE_strategy', 1, allowed=[1,2], fortran_name="sde_strat", comment="decide how Multi-channel should behaves \"1\" means full single diagram enhanced (hep-ph/0208156), \"2\" use the product of the denominator") self.add_param('global_flag', '-O', include=False, hidden=True, comment='global fortran compilation flag, suggestion -fbound-check', fct_mod=(self.make_clean, ('Source'),{})) From 7614a93d5ecc17f86736d4eefd554268fc152d7a Mon Sep 17 00:00:00 2001 From: Olivier Mattelaer Date: Thu, 6 Aug 2026 14:05:56 +0200 Subject: [PATCH 30/42] mg7: keep the flavors carried by processes mapped onto one matrix element p p > l+ l- came out at ~538 pb with apply_flavor_grouping=False against ~1336 pb with it on. The cause is not broken_sym (l+ and l- are distinct, so no identical-particle factor applies) but a missed source of flavor multiplicity in the mg7 exporter. A matrix element carries flavor multiplicity in two independent places: * merged legs (pdg 81/82/...), enumerated by get_external_flavors_with_iden -- the apply_flavor_grouping=True case; * the several processes mapped onto one matrix element when their matrix elements are identical -- the grouping-off case, where u u~ > e+ e-, u u~ > mu+ mu-, c c~ > e+ e- and c c~ > mu+ mu- all share a single matrix element. madevent walks both in get_leshouche_lines, and its leshouche.inc lists all 16 channels of p p > l+ l-. export_mg7 built its flavor list from get_external_flavors_with_iden alone, so with grouping off -- where all the multiplicity sits in the processes list -- subprocesses.json recorded one channel per subprocess instead of four: c/s dropped from the initial state and mu+ mu- dropped entirely. The exporter was already inconsistent with itself, since pdg_color_types in the same function does loop over all processes and duly listed +-4 and +-13. Add HelasMatrixElement.get_flavor_pdg_combinations implementing madevent's rule once, have get_leshouche_lines consume it (behaviour preserving, its IDUP numbering is untouched), and expand the mg7 flavor list through it. Verified: - ungrouped p p > l+ l- now 1335.30 +- 1.65 against grouped 1336.14 +- 1.36, i.e. 0.28 sigma; subprocesses.json carries all 16 channels, structured as the grouped output already was; - madevent output byte-identical across all four apply_flavor_grouping/group_subprocesses combinations (full SubProcesses tree diff, not just leshouche.inc); - 46/46 IOTests pass; - test_flavor_grouping_consistency_mg7, test_generation_from_file_1_mg7, test_group_subprocess_mg7, test_e_e_collision_mg7 all pass. Also add the acceptancetest_mg7_flavor_grouping CI job, which runs test_flavor_grouping_consistency_mg7 -- previously not covered by any workflow, which is why this stayed unnoticed. Co-Authored-By: Claude Opus 5 --- .github/workflows/acceptancetest_mg7.yml | 25 ++++++++++++ madgraph/core/helas_objects.py | 45 +++++++++++++++++++++ madgraph/iolibs/export_mg7.py | 33 +++++++++++++++ madgraph/iolibs/export_v4.py | 32 +++++++-------- tests/acceptance_tests/test_cmd_madevent.py | 17 ++++---- 5 files changed, 126 insertions(+), 26 deletions(-) diff --git a/.github/workflows/acceptancetest_mg7.yml b/.github/workflows/acceptancetest_mg7.yml index 7e8f847276..fe8579fe7b 100644 --- a/.github/workflows/acceptancetest_mg7.yml +++ b/.github/workflows/acceptancetest_mg7.yml @@ -201,6 +201,31 @@ jobs: export PATH="$HOME/.cache/HEPtools/bin:$PATH" ./tests/test_manager.py test_group_subprocess_mg7 -pA -t0 -l INFO + acceptancetest_mg7_flavor_grouping: + needs: build_madspace + # p p > l+ l- must give the same cross-section for all four + # apply_flavor_grouping/group_subprocesses combinations. It used not to: + # with grouping off the mg7 exporter enumerated only the merged-leg flavors + # and dropped the flavors carried by the *processes* mapped onto the same + # matrix element, so 4 of the 16 channels survived (~538 pb instead of + # ~1336 pb). The madevent counterpart is test_flavor_grouping_consistency + # in acceptancetest_madevent.yml. Self-skips if the madspace + + # LHAPDF(NNPDF23) stack is absent. + runs-on: ubuntu-24.04 + if: github.event_name == 'push' || github.event.pull_request.head.repo.fork == true + steps: + - uses: actions/checkout@v4 + - uses: ./.github/actions/checkout_mg5 + - uses: ./.github/actions/install_madspace + - uses: ./.github/actions/restore-pip-cache + - uses: ./.github/actions/restore_heptools + - name: test one of the test test_flavor_grouping_consistency_mg7 + run: | + cd $GITHUB_WORKSPACE + # make the cached lhapdf-config resolvable so the test finds the data dir + export PATH="$HOME/.cache/HEPtools/bin:$PATH" + ./tests/test_manager.py test_flavor_grouping_consistency_mg7 -pA -t0 -l INFO + acceptancetest_mg7_merged_flavor_uq: needs: build_madspace # mg7 cross-section for the merged-flavor u q > u q (q = u d), pinned to the diff --git a/madgraph/core/helas_objects.py b/madgraph/core/helas_objects.py index 46d3b015e7..93a5c43945 100755 --- a/madgraph/core/helas_objects.py +++ b/madgraph/core/helas_objects.py @@ -5680,6 +5680,51 @@ def get_external_flavors_with_iden(self, return_pdgs=False): else: return self['allowed_flavors_with_iden'] + def get_flavor_pdg_combinations(self, model=None): + """Return every physical external-PDG combination this matrix element + covers, grouped per mapped process. + + A matrix element carries flavor multiplicity in *two* independent + places, and both have to be walked to recover all the channels: + + * several `processes` can be mapped onto one matrix element when + their matrix elements are identical. This is what happens with + apply_flavor_grouping=False: u u~ > e+ e-, u u~ > mu+ mu-, + c c~ > e+ e- and c c~ > mu+ mu- all share a single matrix element, + and only the first of them is reachable from its legs; + * a single process can carry *merged* legs (apply_flavor_grouping= + True, pdg 81/82/...), whose concrete flavors are enumerated by + get_external_flavors_with_iden. + + Returns one (pdg_lists, has_merged_particles) pair per process, so that + callers which need the per-process split (madevent's IDUP numbering) + keep it while callers which just want every channel can flatten it. + """ + if model is None: + model = self.get('processes')[0].get('model') + merged = {} + if model and 'merged_particles' in model: + merged = model['merged_particles'] + + combinations = [] + for proc in self.get('processes'): + base_ids = [l.get('id') for l in proc.get_legs_with_decays()] + has_merged = any(abs(pdg) in merged for pdg in base_ids) + if has_merged: + pdg_lists = [] + for flavor in sum(self.get_external_flavors_with_iden(), []): + ids = list(base_ids) + for i, pdg in enumerate(base_ids): + if pdg in merged: + ids[i] = flavor[i] + if -pdg in merged: + ids[i] = -flavor[i] + pdg_lists.append(ids) + else: + pdg_lists = [list(base_ids)] + combinations.append((pdg_lists, has_merged)) + return combinations + def check_flavor_for_all_diagrams(self, real_pdgs, model, debug=False): """Populate every diagram's flavor store for the flavor `real_pdgs`. diff --git a/madgraph/iolibs/export_mg7.py b/madgraph/iolibs/export_mg7.py index 5ca4b9e287..a6cc4f0518 100644 --- a/madgraph/iolibs/export_mg7.py +++ b/madgraph/iolibs/export_mg7.py @@ -19,6 +19,9 @@ def __init__(self, matrix_element, cpp_helas_call_writer): self.diagrams = self.amplitude.get("diagrams") self.helas_diagrams = self.matrix_element.get("diagrams") self.all_flavors, self.all_flavors_pdgs = self.matrix_element.get_external_flavors_with_iden(return_pdgs=True) + self.all_flavors = [list(flavors) for flavors in self.all_flavors] + self.all_flavors_pdgs = [list(pdgs) for pdgs in self.all_flavors_pdgs] + self.expand_flavors_over_processes() self.process = self.amplitude.get("process") self.legs = self.process.get("legs_with_decays") self.color_basis = self.matrix_element.get("color_basis") @@ -43,6 +46,36 @@ def set_topology(self): self.edge_names[number] = f"i{number - 1}" self.incoming[number - 1] = leg.get("id") + def expand_flavors_over_processes(self): + """Add the flavors that live in the *processes* mapped onto this matrix + element rather than in its merged legs. + + get_external_flavors_with_iden only expands merged legs (pdg 81/82/...), + i.e. the apply_flavor_grouping=True case. With grouping off there are no + merged legs and MG5 instead maps every flavor-equivalent process onto a + single matrix element -- u u~ > e+ e-, u u~ > mu+ mu-, c c~ > e+ e- and + c c~ > mu+ mu- all share one -- so asking only for the merged expansion + returns the representative alone and the other channels never make it + into subprocesses.json (p p > l+ l- came out at 538 pb instead of + 1336 pb, i.e. 4 of 16 channels). madevent walks both sources in + get_leshouche_lines; use the same shared enumeration here. + """ + combinations = self.matrix_element.get_flavor_pdg_combinations(self.model) + # Merged legs: get_external_flavors_with_iden already enumerated + # everything, and re-expanding here would double count. + if any(has_merged for _, has_merged in combinations): + return + pdg_lists = [pdgs for pdg_lists, _ in combinations for pdgs in pdg_lists] + if len(pdg_lists) <= 1: + return + # Without merged legs every leg trivially takes flavor index 1, so all + # these processes share the single coupling class and its flavor-index + # tuple; keep all_flavors aligned with all_flavors_pdgs. + if len(self.all_flavors) != 1 or len(self.all_flavors[0]) != 1: + return + self.all_flavors_pdgs = [pdg_lists] + self.all_flavors = [self.all_flavors[0] * len(pdg_lists)] + def set_flavor_indices(self): self.all_flavors_same_initial = [] self.all_flavors_indices = [] diff --git a/madgraph/iolibs/export_v4.py b/madgraph/iolibs/export_v4.py index 337534ff39..36bd3979e2 100755 --- a/madgraph/iolibs/export_v4.py +++ b/madgraph/iolibs/export_v4.py @@ -1728,31 +1728,27 @@ def get_leshouche_lines(self, matrix_element, numproc): lines = [] real_iproc = -1 - for iproc, proc in enumerate(matrix_element.get('processes')): - real_iproc += 1 + # Both sources of flavor multiplicity (several processes mapped onto one + # matrix element, and merged legs within a process) are enumerated by + # HelasMatrixElement.get_flavor_pdg_combinations, shared with the mg7 + # exporter so the two backends cannot drift apart. + processes = matrix_element.get('processes') + for iproc, (pdg_lists, has_merged_particles) in enumerate( + matrix_element.get_flavor_pdg_combinations(self.model)): + proc = processes[iproc] legs = proc.get_legs_with_decays() - ids = [l.get('id') for l in legs] - has_merged_particles = False - if self.model and 'merged_particles' in self.model: - has_merged_particles = any([abs(id) in self.model['merged_particles'] for id in ids]) + real_iproc += 1 if has_merged_particles: - allow_flavor = matrix_element.get_external_flavors_with_iden() - for flavor in sum(allow_flavor,[]): - ids = [l.get('id') for l in legs] - for i,id in enumerate(ids): - if id in self.model['merged_particles']: - ids[i] = flavor[i] #self.model['merged_particles'][id][flavor[i]-1] - if -id in self.model['merged_particles']: - ids[i] = -flavor[i] #self.model['merged_particles'][-id][flavor[i]-1] + for ids in pdg_lists: lines.append("DATA (IDUP(i,%d,%d),i=1,%d)/%s/" % \ (real_iproc + 1, numproc+1, nexternal, - ",".join([str(id) for id in ids]))) - real_iproc += 1 + ",".join([str(id) for id in ids]))) + real_iproc += 1 else: lines.append("DATA (IDUP(i,%d,%d),i=1,%d)/%s/" % \ (real_iproc + 1, numproc+1, nexternal, - ",".join([str(l.get('id')) for l in legs]))) - + ",".join([str(id) for id in pdg_lists[0]]))) + if iproc == 0 and numproc == 0: for i in [1, 2]: lines.append("DATA (MOTHUP(%d,i),i=1,%2r)/%s/" % \ diff --git a/tests/acceptance_tests/test_cmd_madevent.py b/tests/acceptance_tests/test_cmd_madevent.py index 6431bdcfe5..da0107c979 100755 --- a/tests/acceptance_tests/test_cmd_madevent.py +++ b/tests/acceptance_tests/test_cmd_madevent.py @@ -1337,14 +1337,15 @@ def test_flavor_grouping_consistency(self): def test_flavor_grouping_consistency_mg7(self): """mg7 equivalent of test_flavor_grouping_consistency for p p > l+ l-. - KNOWN-FAILING, intentionally NOT marked xfail: mg7 currently returns - cross-sections that depend on the apply_flavor_grouping setting (e.g. - ~1332 pb grouped vs ~511 pb ungrouped) because the broken-symmetry - (flavour-consolidation) factor implemented for standalone / - standalone_cpp is not yet applied on the mg7 (madmatrix) side. The four - settings must agree; the test asserts that and is left undecorated so the - mg7 flavour-grouping discrepancy stays visible until broken_sym is ported - to mg7. It self-skips where the mg7 runtime stack is unavailable. + The four settings must all give the same cross-section. They used not + to (~1336 pb grouped vs ~538 pb ungrouped): a matrix element carries + flavor multiplicity either in its merged legs (apply_flavor_grouping= + True) or in the several processes mapped onto it (grouping off), and the + mg7 exporter only enumerated the first, so with grouping off it kept 4 + of the 16 channels of p p > l+ l- -- dropping c/s in the initial state + and mu+ mu- entirely. Both sources are now walked through the shared + HelasMatrixElement.get_flavor_pdg_combinations. It self-skips where the + mg7 runtime stack is unavailable. """ datadir = _mg7_datadir_or_skip(self) settings = [ From eca43df178a549b9d1f83b18295633ce5c517514 Mon Sep 17 00:00:00 2001 From: Olivier Mattelaer Date: Thu, 6 Aug 2026 17:14:12 +0200 Subject: [PATCH 31/42] make a failed run fail as a failed run, not as a disagreement test_flavor_grouping_consistency intermittently reported "Incompatible cross-sections ... (1e-99 +- 0)", which reads like a physics disagreement but means the run produced nothing at all. Two things conspired to hide that. First, the guards could not catch it. Each configuration checked its own precision with err/(cross+1e-99); when a run fails, cross and error are both 0, so that is 0/1e-99 = 0 < 0.05 and it passes. The explicit zero-checks were dead code for the same reason: two of them tested val == 0 *after* adding 1e-99, and the MLM one asserted val > 0 on the padded value. So a zero always slipped through to the pairwise comparison, where it produced a misleading message. Assert on the raw cross instead, in all four affected tests (the three madevent ones and the mg7 counterpart), naming the run directory. Second, when such a run does fail the cause is usually a ZeroResult swallowed by nice_error_handling, which only logs a warning -- invisible at the CRITICAL level the tests run at -- leaving cross/error at 0. The new assertion turns that from a puzzle into a one-line diagnosis. Also fix multiple_try: after exhausting its retries it ended with a bare `raise` outside the except block, so instead of re-raising the real error it raised "RuntimeError: No active exception to reraise", discarding the only information explaining the failure (and doing so after several sleeps). Hit while debugging the above, where it masked an AttributeError from load_results_db. Raise my_error, which keeps the traceback; the non-__debug__ branch is unchanged. Verified: the new guard fires on cross=0 and stays silent on a good run; multiple_try now surfaces the original exception with its traceback under __debug__ and the wrapped message under -O. test_flavor_grouping_ consistency, _width, _mg7 and test_generation_from_file_1_mg7 all pass. Note this makes a future occurrence self-diagnosing; it does not fix the underlying transient, which I could not reproduce (~28 configurations, idle and under load). A likely contributor is gen_ximprove.get_helicity classifying a subprocess as having no phase space purely from `if stdout:` on ./gensym, without ever checking its return code. Co-Authored-By: Claude Opus 5 --- madgraph/various/misc.py | 8 ++- tests/acceptance_tests/test_cmd_madevent.py | 68 ++++++++++++++++----- 2 files changed, 60 insertions(+), 16 deletions(-) diff --git a/madgraph/various/misc.py b/madgraph/various/misc.py index ddac602ac0..15e28d903b 100755 --- a/madgraph/various/misc.py +++ b/madgraph/various/misc.py @@ -473,7 +473,13 @@ def deco_f_retry(*args, **opt): time.sleep(sleep * (i+1)) if __debug__: - raise + # Re-raise the original error, traceback included. A bare + # `raise` here is *outside* the except block, so it does not + # re-raise my_error: it fails with "RuntimeError: No active + # exception to reraise" and hides the real cause (which, after + # nb_try attempts and their sleeps, is the only thing that + # explains what went wrong). + raise my_error raise my_error.__class__('[Fail %i times] \n %s ' % (i+1, my_error)) return deco_f_retry return deco_retry diff --git a/tests/acceptance_tests/test_cmd_madevent.py b/tests/acceptance_tests/test_cmd_madevent.py index da0107c979..e779bacac5 100755 --- a/tests/acceptance_tests/test_cmd_madevent.py +++ b/tests/acceptance_tests/test_cmd_madevent.py @@ -1297,13 +1297,25 @@ def test_flavor_grouping_consistency(self): self.do('generate_events -f') - val = self.cmd_line.results.current['cross'] + 1e-99 + cross = self.cmd_line.results.current['cross'] err = self.cmd_line.results.current['error'] - results.append((val, err, afg, gsp)) - if val == 0: - misc.sprint('Warning: cross-section is zero for ' - 'apply_flavor_grouping=%s/group_subprocesses=%s' % (afg, gsp)) + # A run that produced nothing leaves cross *and* error at 0 -- most + # often a ZeroResult swallowed by nice_error_handling, which only + # logs a warning (invisible at the CRITICAL level these tests run + # at). Catch it here: the precision check below divides by + # cross+1e-99, so 0/1e-99 = 0 < 0.05 sails through, and the zero + # would surface only in the pairwise comparison as a misleading + # "incompatible cross-sections ... (1e-99 +- 0)" rather than as the + # failed run it is. (The check this replaces tested val == 0 after + # the +1e-99, so it could never fire.) + self.assertTrue(cross, + 'no cross-section produced for apply_flavor_grouping=%s/' + 'group_subprocesses=%s: the run failed rather than disagreeing ' + '(cross=%s, error=%s, run dir %s)' % (afg, gsp, cross, err, run_dir)) + + val = cross + 1e-99 + results.append((val, err, afg, gsp)) #check precision is reasonable for each individual run self.assertLess(err / val, 0.05, @@ -1364,6 +1376,13 @@ def test_flavor_grouping_consistency_mg7(self): 'set group_subprocesses %s' % gsp, 'generate p p > l+ l-'], pjoin(self.path, 'MG7_fg_%d' % i), datadir) + # Fail on a run that produced nothing rather than letting it reach + # the pairwise comparison: err/(cross+1e-99) is 0 < 0.05 when both + # are 0, so the precision check below cannot catch it. + self.assertTrue(cross, + 'mg7 produced no cross-section for apply_flavor_grouping=%s/' + 'group_subprocesses=%s: the run failed rather than disagreeing ' + '(cross=%s, error=%s)' % (afg, gsp, cross, err)) results.append((cross + 1e-99, err, afg, gsp)) self.assertLess(err / (cross + 1e-99), 0.05, 'mg7 cross-section too imprecise (afg=%s, gsp=%s): %s +- %s' @@ -1419,13 +1438,25 @@ def test_flavor_grouping_consistency_width(self): self.do('generate_events -f') - val = self.cmd_line.results.current['cross'] + 1e-99 + cross = self.cmd_line.results.current['cross'] err = self.cmd_line.results.current['error'] - results.append((val, err, afg, gsp)) - if val == 0: - misc.sprint('Warning: cross-section is zero for ' - 'apply_flavor_grouping=%s/group_subprocesses=%s' % (afg, gsp)) + # A run that produced nothing leaves cross *and* error at 0 -- most + # often a ZeroResult swallowed by nice_error_handling, which only + # logs a warning (invisible at the CRITICAL level these tests run + # at). Catch it here: the precision check below divides by + # cross+1e-99, so 0/1e-99 = 0 < 0.05 sails through, and the zero + # would surface only in the pairwise comparison as a misleading + # "incompatible cross-sections ... (1e-99 +- 0)" rather than as the + # failed run it is. (The check this replaces tested val == 0 after + # the +1e-99, so it could never fire.) + self.assertTrue(cross, + 'no cross-section produced for apply_flavor_grouping=%s/' + 'group_subprocesses=%s: the run failed rather than disagreeing ' + '(cross=%s, error=%s, run dir %s)' % (afg, gsp, cross, err, run_dir)) + + val = cross + 1e-99 + results.append((val, err, afg, gsp)) #check precision is reasonable for each individual run self.assertLess(err / val, 0.05, @@ -1514,13 +1545,20 @@ def test_flavor_grouping_consistency_mlm(self): self.do('generate_events -f') # Verify event generation succeeded - val = self.cmd_line.results.current['cross'] + 1e-99 + cross = self.cmd_line.results.current['cross'] err = self.cmd_line.results.current['error'] - results.append((val, err, afg, gsp)) - # Check that we got a valid cross-section - self.assertGreater(val, 0, - 'cross-section is zero for q q~ > q q~ with MLM merging') + # Check that we got a valid cross-section. Test the raw cross, not + # cross+1e-99: a run that produced nothing leaves cross and error at + # 0 (typically a ZeroResult swallowed by nice_error_handling), and + # asserting on the padded value can never fail. + self.assertTrue(cross, + 'no cross-section produced for q q~ > q q~ with MLM merging: ' + 'the run failed (cross=%s, error=%s, run dir %s)' + % (cross, err, run_dir)) + + val = cross + 1e-99 + results.append((val, err, afg, gsp)) # Check precision is reasonable self.assertLess(err / val, 0.10, From ab161ac8a7099d2e3824f1c96752a58e2e29df98 Mon Sep 17 00:00:00 2001 From: Olivier Mattelaer Date: Fri, 7 Aug 2026 18:32:21 +0200 Subject: [PATCH 32/42] build the (n-2)! Del Duca-Dixon-Maltoni basis for multi-gluon processes For a process whose colour structure is purely adjoint, the Jacobi identity lets any colour factor be written on the (n-2)! half-ladder structures instead of the (n-1)! traces. The reduction walks each diagram's colour string as a tree and expands every subtree hanging off the leg-1-to-leg-n spine as a nested commutator of adjoint generators, which costs 2^(n-2) per diagram. Selected by 'set color_basis auto|trace|ddm', with auto turning it on for the exporters which declare support_ddm_color_basis. Fortran standalone and the grouped madevent exporter do; everything else keeps the trace basis, and any process which is not fully adjoint falls back on its own. Two things this needed: - f.complex_conjugate now returns f untouched. The inherited implementation reverses the indices, which is right for Tr but flips the sign of the totally antisymmetric f. It was latent because full_simplify removes every f before the trace-basis colour matrix is built; with DDM the basis elements are products of (n-2) f's, so |M|^2 came out negative for odd n. - ColorMatrix.build_matrix_ddm expands each ladder once onto its 2^(n-2) traces and assembles the entry from cached trace-trace products. Contracting two ladders head on turns every one of the 2(n-2) f's into a pair of traces: gg>gggg spent 17.4s there against 0.089s in the trace basis, and gg>ggggg never finished. It is now 0.12s and 1.6s. madevent keeps the trace basis for the colour flows: the colour sum runs on the DDM structures while JAMP2 is filled from trace JAMPs rebuilt out of the DDM ones through the Kleiss-Kuijf relations, 3840 terms instead of 231840 at n=7. Co-Authored-By: Claude Opus 5 --- madgraph/core/color_algebra.py | 7 + madgraph/core/color_amp.py | 542 +++++++++++++++++- madgraph/interface/madgraph_interface.py | 72 ++- madgraph/interface/master_interface.py | 3 + madgraph/iolibs/export_v4.py | 86 ++- .../matrix_madevent_group_v4.inc | 10 +- .../matrix_madevent_group_v4_hel.inc | 10 +- tests/unit_tests/core/test_color_algebra.py | 5 + tests/unit_tests/core/test_color_amp.py | 131 +++++ 9 files changed, 845 insertions(+), 21 deletions(-) diff --git a/madgraph/core/color_algebra.py b/madgraph/core/color_algebra.py index 75fc891444..662f124f28 100755 --- a/madgraph/core/color_algebra.py +++ b/madgraph/core/color_algebra.py @@ -332,6 +332,13 @@ def simplify(self): return ColorFactor([col_str1, col_str2]) + def complex_conjugate(self): + """f (and d) are real, so complex conjugation leaves them untouched. + The default behaviour of reversing the indices would introduce a + spurious sign, since f is totally antisymmetric.""" + + return self + #=============================================================================== # d #=============================================================================== diff --git a/madgraph/core/color_amp.py b/madgraph/core/color_amp.py index 04a27c6ee2..0437ab2684 100755 --- a/madgraph/core/color_amp.py +++ b/madgraph/core/color_amp.py @@ -22,6 +22,8 @@ import copy import fractions import itertools +import itertools +import logging import operator import re import array @@ -36,6 +38,242 @@ if madgraph.ordering: set = misc.OrderedSet +logger = logging.getLogger('madgraph.color_amp') + +#=============================================================================== +# Del Duca-Dixon-Maltoni (adjoint) color basis +#=============================================================================== +# For a process whose color structure is purely adjoint (all colored external +# legs are octets) the historical trace basis Tr(1,sigma(2),...,sigma(n)) has +# (n-1)! elements while, thanks to the Jacobi identity, the color factor of any +# such amplitude can be written on the (n-2)! "half-ladder" (multi-peripheral) +# structures +# F(sigma) = f(1,sigma(2),x1) f(x1,sigma(3),x2) ... f(x(n-3),sigma(n-1),n) +# where legs 1 and n are kept fixed at the two ends of the ladder. This is the +# Del Duca-Dixon-Maltoni basis. Using it divides the number of JAMPs by (n-1) +# and the size of the color matrix by (n-1)^2. +# +# Module level switch selecting the color basis used for fully adjoint +# processes. Set through 'set color_basis' in the MG5 interface (the exporters +# which need a color flow decomposition, i.e. anything writing leshouche +# information, must keep the trace basis). +ddm_basis = False +# Whether the trace basis must be built next to the DDM one. Needed by the +# output formats which have to assign a color flow to an event: the color sum +# then runs over the (n-2)! DDM structures while the color flow probabilities +# keep using the (n-1)! trace ones, obtained from the DDM JAMPs through the +# Kleiss-Kuijf relations. +ddm_flow_basis = False + + +def set_ddm_basis(value, with_flow=False): + """Set the module wide switch selecting the DDM color basis.""" + + global ddm_basis, ddm_flow_basis + ddm_basis = bool(value) + ddm_flow_basis = ddm_basis and bool(with_flow) + + +class DDMError(Exception): + """Raised when a color string cannot be mapped onto the DDM basis. Always + caught by ColorBasis.build, which then falls back on the trace basis.""" + + +def ddm_half_ladder(perm, first, last): + """Return the immutable color string of the DDM half-ladder structure + f(first,perm[0],-1) f(-1,perm[1],-2) ... f(-(m-1),perm[m-1],last) + for the ordered tuple perm of the (n-2) legs sitting between the two fixed + ends first and last.""" + + if len(perm) == 1: + col_objs = [color_algebra.f(first, perm[0], last)] + else: + col_objs = [color_algebra.f(first, perm[0], -1)] + col_objs.extend([color_algebra.f(-(i + 1), leg, -(i + 2)) \ + for i, leg in enumerate(perm[1:-1])]) + col_objs.append(color_algebra.f(-(len(perm) - 1), perm[-1], last)) + + return color_algebra.ColorString(col_objs).to_immutable() + + +def _reorder_sign(stored, wanted): + """Signature of the permutation bringing the three indices of an f object + from the order 'stored' to the order 'wanted'. f is totally antisymmetric, + so f(stored) = _reorder_sign(stored,wanted) * f(wanted).""" + + perm = [stored.index(index) for index in wanted] + sign = 1 + for i in range(len(perm)): + for j in range(i + 1, len(perm)): + if perm[i] > perm[j]: + sign = -sign + + return sign + + +class _ColorTree(object): + """A product of f objects seen as a tree: the f's are the nodes, the summed + (negative) indices the internal edges and the external (positive) indices + the leaves. Provides the reduction onto the DDM half-ladder basis.""" + + def __init__(self, col_str): + """Build the tree from a ColorString made of f objects only. Raise + DDMError as soon as the string is not a fully adjoint color tree.""" + + self.nodes = [] + for col_obj in col_str: + if col_obj.__class__.__name__ == 'ColorOne': + continue + if type(col_obj) is not color_algebra.f: + raise DDMError("%s is not an f object" % str(col_obj)) + self.nodes.append(tuple(col_obj)) + + if not self.nodes: + raise DDMError("empty color string") + + # Locate each index. External indices must appear once, summed ones + # exactly twice and in two different nodes. + self.where = collections.defaultdict(list) + for i, node in enumerate(self.nodes): + if node[0] == node[1] or node[1] == node[2] or node[0] == node[2]: + raise DDMError("f object %s has repeated indices" % str(node)) + for index in node: + self.where[index].append(i) + + self.externals = [] + nb_internal = 0 + for index, nodes in self.where.items(): + if index > 0: + if len(nodes) != 1: + raise DDMError("external index %i appears %i times" % \ + (index, len(nodes))) + self.externals.append(index) + else: + if len(nodes) != 2: + raise DDMError("summed index %i appears %i times" % \ + (index, len(nodes))) + nb_internal += 1 + + # A connected graph with V nodes and V-1 edges is a tree + if nb_internal != len(self.nodes) - 1: + raise DDMError("color structure is not a tree") + + def _neighbour(self, index, node): + """The node sharing the summed index 'index' with node 'node'.""" + + first, second = self.where[index] + return second if first == node else first + + def _spine(self, first, last): + """The list of nodes on the path going from the leaf 'first' to the + leaf 'last'.""" + + if first not in self.externals or last not in self.externals: + raise DDMError("legs %s and %s are not both external here" % \ + (first, last)) + start = self.where[first][0] + end = self.where[last][0] + + # Depth first search on the node tree, keeping track of the path + stack = [(start, None, [start])] + while stack: + node, from_index, path = stack.pop() + if node == end: + return path + for index in self.nodes[node]: + if index > 0 or index == from_index: + continue + stack.append((self._neighbour(index, node), index, + path + [self._neighbour(index, node)])) + + raise DDMError("color structure is not connected") + + def _subtree(self, index, from_node): + """Expansion of the adjoint matrix associated to the subtree hanging on + the edge 'index' of node 'from_node', as a list of + (sign, ordered tuple of legs). A single leaf gives the generator + itself, and a node with two children A and B gives the commutator + [M_B, M_A] (which is where the (n-2)! counting comes from).""" + + if index > 0: + return [(1, (index,))] + + node_i = self._neighbour(index, from_node) + node = self.nodes[node_i] + children = list(node) + children.remove(index) + alpha, beta = children + sign = _reorder_sign(node, (index, alpha, beta)) + + exp_a = self._subtree(alpha, node_i) + exp_b = self._subtree(beta, node_i) + + result = [] + for (sa, wa), (sb, wb) in itertools.product(exp_a, exp_b): + result.append((sign * sa * sb, wb + wa)) + result.append((-sign * sa * sb, wa + wb)) + + return result + + def reduce_to_ddm(self, first, last): + """Decompose the tree onto the DDM half-ladder basis with the legs + 'first' and 'last' at the two ends. Returns {ordered legs: coefficient} + where the keys are the (n-2) other external legs in ladder order.""" + + spine = self._spine(first, last) + + # Split each node of the spine into (incoming, hanging, outgoing) + global_sign = 1 + hanging = [] + for pos, node_i in enumerate(spine): + node = self.nodes[node_i] + if pos == 0: + in_index = first + else: + in_index = [i for i in node if i in self.nodes[spine[pos - 1]]][0] + if pos == len(spine) - 1: + out_index = last + else: + out_index = [i for i in node if i in self.nodes[spine[pos + 1]]][0] + off_index = [i for i in node if i not in (in_index, out_index)][0] + global_sign *= _reorder_sign(node, (in_index, off_index, out_index)) + hanging.append(self._subtree(off_index, node_i)) + + # The whole structure is the matrix product M_m ... M_1 between the + # ends, and (M_i1 ... M_ik) contracted between 'last' and 'first' is the + # half-ladder with the legs in the reversed order. + result = collections.defaultdict(int) + for combination in itertools.product(*reversed(hanging)): + sign = global_sign + word = [] + for term_sign, term_word in combination: + sign *= term_sign + word.extend(term_word) + result[tuple(reversed(word))] += sign + + return dict((perm, coeff) for perm, coeff in result.items() if coeff) + + +def reduce_to_ddm(col_str, first, last): + """Decompose the ColorString col_str (a product of f objects) onto the DDM + half-ladder basis, returning a ColorFactor whose strings are the basis + elements. Raise DDMError if col_str is not a fully adjoint color tree.""" + + decomposition = _ColorTree(col_str).reduce_to_ddm(first, last) + + col_fact = color_algebra.ColorFactor() + for perm, coeff in decomposition.items(): + new_str = color_algebra.ColorString() + new_str.from_immutable(ddm_half_ladder(perm, first, last)) + new_str.coeff = col_str.coeff * coeff + new_str.is_imaginary = col_str.is_imaginary + new_str.Nc_power = col_str.Nc_power + new_str.loop_Nc_power = col_str.loop_Nc_power + col_fact.append(new_str) + + return col_fact + + #=============================================================================== # ColorBasis #=============================================================================== @@ -61,6 +299,15 @@ class ColorBasis(dict): # permute_immutable (Tr is cyclic, T is an open chain, ColorOne is empty). fast_relabel_objects = frozenset(['Tr', 'T', 'ColorOne']) + # Legs at the two ends of the DDM half-ladders (None for the trace basis) + _ddm_ends = None + + # Trace basis built next to a DDM one, carrying the color flows + _flow_basis = None + + # Dictionary to save the DDM decompositions already done + _ddm_dict = {} + class ColorBasisError(Exception): """Exception raised if an error occurs in the definition @@ -305,8 +552,128 @@ def relabel_canonical(self, col_fact, canonical_rep): self._fast_relabel_dict[canonical_rep] = verdict return fast if verdict else slow + def get_ddm_ends(self): + """If every color structure of this basis is a fully adjoint color tree + over one and the same set of external legs -- i.e. if the process is a + pure multi-gluon one -- return the two legs to be put at the ends of the + DDM half-ladders. Return None otherwise.""" + + legs = None + for colorize_dict in self._list_color_dict: + for col_str in colorize_dict.values(): + externals = [] + for col_obj in col_str: + if col_obj.__class__.__name__ == 'ColorOne': + continue + if type(col_obj) is not color_algebra.f: + return None + externals.extend([i for i in col_obj if i > 0]) + if len(externals) != len(set(externals)): + return None + externals = sorted(externals) + if legs is None: + legs = externals + if len(legs) < 3: + return None + elif legs != externals: + return None + + if legs is None: + return None + + return (legs[0], legs[-1]) + + def update_color_basis_ddm(self, colorize_dict, index): + """Same as update_color_basis, but decomposing the color structures on + the (n-2)! DDM half-ladder basis instead of the (n-1)! trace one.""" + + first, last = self._ddm_ends + + for col_chain, col_str in colorize_dict.items(): + # The decomposition only depends on the tree structure, so + # normalize the summed indices to make the cache hit as often as + # possible. + repl_dict = {} + for col_obj in col_str: + for i in col_obj: + if i < 0 and i not in repl_dict: + repl_dict[i] = -len(repl_dict) - 1 + canonical_str = col_str.create_copy() + canonical_str.replace_indices(repl_dict) + canonical_rep = canonical_str.to_immutable() + + try: + decomposition = self._ddm_dict[canonical_rep] + except KeyError: + decomposition = _ColorTree(canonical_str).reduce_to_ddm(first, + last) + self._ddm_dict[canonical_rep] = decomposition + + for perm, coeff in decomposition.items(): + basis_entry = (index, + col_chain, + col_str.coeff * coeff, + col_str.is_imaginary, + col_str.Nc_power, + col_str.loop_Nc_power) + immutable_col_str = ddm_half_ladder(perm, first, last) + try: + self[immutable_col_str].append(basis_entry) + except KeyError: + self[immutable_col_str] = [basis_entry] + + def build_flow_basis(self): + """Build, next to the DDM basis, the trace basis which is the one + carrying the color flow information. Only the basis is built, not its + (n-1)!^2 color matrix, since the color sum stays in the DDM basis.""" + + flow_basis = ColorBasis() + flow_basis._list_color_dict = self._list_color_dict + for index, color_dict in enumerate(self._list_color_dict): + flow_basis.update_color_basis(color_dict, index) + + self._flow_basis = flow_basis + + def get_flow_basis(self): + """The color basis carrying the color flow information: the trace basis + built next to the DDM one, or simply self for a trace basis.""" + + return self._flow_basis if self._flow_basis else self + + def get_flow_projection(self): + """Return the Kleiss-Kuijf relations giving each trace JAMP as a linear + combination of the DDM ones, i.e. the coefficients of the expansion of + the half-ladders on the trace basis, transposed. The format is the one + of get_color_amplitudes, so that the same writers can be used: a list + (one entry per element of the flow basis) of + ((1, coefficient, is_imaginary, Nc power), DDM basis index+1).""" + + if not self._flow_basis: + raise ColorBasis.ColorBasisError( + "No flow basis attached to this color basis") + + flow_index = dict((struct, i) for i, struct in \ + enumerate(sorted(self._flow_basis.keys()))) + projection = [[] for i in range(len(flow_index))] + + for i, struct in enumerate(sorted(self.keys())): + col_str = color_algebra.ColorString() + col_str.from_immutable(struct) + for cs in color_algebra.ColorFactor([col_str]).full_simplify(): + try: + row = flow_index[cs.to_immutable()] + except KeyError: + raise ColorBasis.ColorBasisError( + "The half-ladder %s expands on the trace structure %s " + "which is not part of the flow basis" % \ + (str(col_str), str(cs))) + projection[row].append(((1, cs.coeff, cs.is_imaginary, + cs.Nc_power), i + 1)) + + return projection + def update_color_basis(self, colorize_dict, index): - """Update the current color basis by adding information from + """Update the current color basis by adding information from the colorize dictionary (produced by the colorize routine) associated to diagram with index index. Keep track of simplification results for maximal optimization.""" @@ -400,17 +767,33 @@ def create_color_dict_list(self, amplitude): def build(self, amplitude=None): """Build the a color basis object using information contained in - amplitude (otherwise use info from _list_color_dict). + amplitude (otherwise use info from _list_color_dict). Returns a list of color """ if amplitude: self.create_color_dict_list(amplitude) + + if ddm_basis: + self._ddm_ends = self.get_ddm_ends() + if self._ddm_ends: + try: + for index, color_dict in enumerate(self._list_color_dict): + self.update_color_basis_ddm(color_dict, index) + if ddm_flow_basis: + self.build_flow_basis() + return + except DDMError as error: + logger.debug('Falling back on the trace color basis: %s', error) + self.clear() + self._ddm_ends = None + self._flow_basis = None + for index, color_dict in enumerate(self._list_color_dict): self.update_color_basis(color_dict, index) def __init__(self, *args): """Initialize a new color basis object, either empty or filled (0 - or 1 arguments). If one arguments is given, it's interpreted as + or 1 arguments). If one arguments is given, it's interpreted as an amplitude.""" assert len(args) < 2, "Object ColorBasis must be initialized with 0 or 1 arguments" @@ -427,6 +810,16 @@ def __init__(self, *args): # Whether relabel_canonical may take its shortcut, per canonical form self._fast_relabel_dict = {} + # Legs at the two ends of the DDM half-ladders, None when the basis is + # the standard trace one + self._ddm_ends = None + + # Trace basis built next to a DDM one, carrying the color flows + self._flow_basis = None + + # Dictionary to save the DDM decompositions already done + self._ddm_dict = {} + if args: assert isinstance(args[0], diagram_generation.Amplitude), \ @@ -529,6 +922,11 @@ def color_flow_decomposition(self, repr_dict, ninitial): here (an error is raised). Needs a dictionary with keys being external leg numbers, and value the corresponding color representation.""" + if self._ddm_ends: + raise ColorBasis.ColorBasisError( + "A DDM color basis has no single color flow per basis element." + " Use 'set color_basis trace' for this output format.") + # Offsets used to introduce fake quark indices for gluons offset1 = 1000 offset2 = 2000 @@ -947,6 +1345,7 @@ class ColorMatrix(dict): _col_basis1 = None _col_basis2 = None col_matrix_fixed_Nc = {} + _ddm_expansions = None def __init__(self, col_basis, col_basis2=None, Nc=3, Nc_power_min=None, Nc_power_max=None): @@ -962,6 +1361,8 @@ def __init__(self, col_basis, col_basis2=None, self._val_index = array.array('i') self._sorted_keys1 = [] self._sorted_keys2 = [] + # Set by setup_ddm_entries for a DDM (half-ladder) color basis + self._ddm_expansions = None self.col_matrix_fixed_Nc = _ColorMatrixView(self, 1) self._col_basis1 = col_basis @@ -1108,6 +1509,10 @@ def build_matrix(self, Nc=3, if not n1 or not n2: return + if getattr(self._col_basis1, '_ddm_ends', None) and \ + getattr(self._col_basis2, '_ddm_ends', None): + self.setup_ddm_entries() + canonical_dict = {} symmetry = ColorBasisSymmetry(keys1, None if keys2 is keys1 else keys2) @@ -1160,12 +1565,141 @@ def build_matrix(self, Nc=3, assert progressed, "Color matrix orbit exploration made no progress" remaining = still_missing + def setup_ddm_entries(self): + """Switch create_new_entry over to the assembly used for a DDM + (half-ladder) color basis. + + Contracting two half-ladders head on is exponentially expensive, since + the simplification rules turn every one of the 2(n-2) f objects into a + pair of traces. Each ladder is instead expanded once on the trace basis + (2^(n-2) traces) and the entry is assembled from trace-trace products, + which are recycled between all the entries. Everything else, including + the orbit symmetry of the basis, is left to build_matrix.""" + + self._ddm_expansions = {} + self._ddm_half_dict = {} + self._ddm_trace_dict = {} + + def get_ddm_trace_expansion(self, struct): + """Expansion of one half-ladder on the trace basis, as a list of + (immutable trace, coefficient, is_imaginary, Nc power).""" + + try: + return self._ddm_expansions[struct] + except KeyError: + pass + + col_str = color_algebra.ColorString() + col_str.from_immutable(struct) + expansion = [(cs.to_immutable(), cs.coeff, cs.is_imaginary, + cs.Nc_power) for cs in \ + color_algebra.ColorFactor([col_str]).full_simplify()] + + self._ddm_expansions[struct] = expansion + return expansion + + def create_new_entry_ddm(self, struct1, struct2, + Nc_power_min, Nc_power_max, Nc): + """create_new_entry for two half-ladders, through their trace + expansions.""" + + contraction = collections.defaultdict(fractions.Fraction) + for trace, coeff, is_imaginary, Nc_power in \ + self.get_ddm_trace_expansion(struct1): + self.accumulate_number(contraction, + (coeff, is_imaginary, Nc_power), + self.get_half_ladder_contraction(trace, + struct2)) + + result = color_algebra.ColorFactor() + for (is_imaginary, Nc_power), coeff in contraction.items(): + if not coeff: + continue + if Nc_power_min is not None and Nc_power < Nc_power_min: + continue + if Nc_power_max is not None and Nc_power > Nc_power_max: + continue + result.append(color_algebra.ColorString([], coeff, is_imaginary, + Nc_power)) + + return result, result.set_Nc(Nc) + + @staticmethod + def accumulate_number(target, factor, numbers): + """Add factor*numbers to target, where a number is a dictionary + {(is_imaginary, Nc power): coefficient} and factor a single + (coefficient, is_imaginary, Nc power) triplet.""" + + coeff, is_imaginary, Nc_power = factor + for (other_imaginary, other_power), other_coeff in numbers.items(): + new_coeff = coeff * other_coeff + if is_imaginary and other_imaginary: + new_coeff = -new_coeff + new_imaginary = False + else: + new_imaginary = is_imaginary or other_imaginary + target[(new_imaginary, Nc_power + other_power)] += new_coeff + + def get_half_ladder_contraction(self, trace, struct2): + """Contraction of the single trace \'trace\' with the complex conjugate + of the half-ladder \'struct2\'.""" + + canonical_rep, dummy = \ + color_algebra.ColorString().to_canonical(trace + struct2) + try: + return self._ddm_half_dict[canonical_rep] + except KeyError: + pass + + result = collections.defaultdict(fractions.Fraction) + for trace2, coeff, is_imaginary, Nc_power in \ + self.get_ddm_trace_expansion(struct2): + # complex conjugation of the coefficient of the second ladder + if is_imaginary: + coeff = -coeff + self.accumulate_number(result, (coeff, is_imaginary, Nc_power), + self.get_trace_contraction(trace, trace2)) + + self._ddm_half_dict[canonical_rep] = result + return result + + def get_trace_contraction(self, trace1, trace2): + """Contraction of two single traces, as a dictionary + {(is_imaginary, Nc power): coefficient}.""" + + canonical_rep, dummy = \ + color_algebra.ColorString().to_canonical(trace1 + trace2) + try: + return self._ddm_trace_dict[canonical_rep] + except KeyError: + pass + + col_str = color_algebra.ColorString() + col_str.from_immutable(trace1) + col_str2 = color_algebra.ColorString() + col_str2.from_immutable(trace2) + col_str.product(col_str2.complex_conjugate()) + + result = collections.defaultdict(fractions.Fraction) + for cs in color_algebra.ColorFactor([col_str]).full_simplify(): + assert not len(cs), \ + "Trace contraction %s did not simplify to a number" % str(cs) + result[(cs.is_imaginary, cs.Nc_power)] += cs.coeff + + self._ddm_trace_dict[canonical_rep] = result + return result + + def create_new_entry(self, struct1, struct2, Nc_power_min, Nc_power_max, Nc): """ Create a new product result, and result with fixed Nc for two color basis entries. Implement Nc power limits.""" - # Create color string objects corresponding to color basis + if self._ddm_expansions is not None: + return self.create_new_entry_ddm(struct1, struct2, + Nc_power_min, Nc_power_max, Nc) + + # Create color string objects corresponding to color basis # keys col_str = color_algebra.ColorString() col_str.from_immutable(struct1) diff --git a/madgraph/interface/madgraph_interface.py b/madgraph/interface/madgraph_interface.py index d7e1ef7fd0..d07744e32c 100755 --- a/madgraph/interface/madgraph_interface.py +++ b/madgraph/interface/madgraph_interface.py @@ -60,6 +60,7 @@ import madgraph.core.base_objects as base_objects +import madgraph.core.color_amp as color_amp import madgraph.core.diagram_generation as diagram_generation import madgraph.loop.loop_diagram_generation as loop_diagram_generation import madgraph.loop.loop_base_objects as loop_base_objects @@ -3155,8 +3156,10 @@ class MadGraphCmd(HelpToCmd, CheckValidForCmd, CompleteForCmd, CmdExtended): 'max_t_for_channel', 'zerowidth_tchannel', 'default_unset_couplings', - 'nlo_mixed_expansion' + 'nlo_mixed_expansion', + 'color_basis' ] + _valid_color_basis = ['auto', 'trace', 'ddm'] _valid_nlo_modes = ['all','real','virt','sqrvirt','tree','noborn','LOonly', 'only'] _valid_sqso_types = ['==','<=','=','>'] _valid_amp_so_types = ['=','<=', '==', '>'] @@ -3238,7 +3241,8 @@ class MadGraphCmd(HelpToCmd, CheckValidForCmd, CompleteForCmd, CmdExtended): 'max_t_for_channel': 99, # means no restrictions 'zerowidth_tchannel': True, 'nlo_mixed_expansion':True, - 'apply_flavor_grouping': True + 'apply_flavor_grouping': True, + 'color_basis': 'auto' } options_madevent = {'automatic_html_opening':True, @@ -4239,6 +4243,9 @@ def create_lambda_values_list(lower_bound, N): ###### BEGIN do_check + # No exporter here, so 'auto' means the safe trace basis + self.set_color_basis_mode() + args = self.split_arg(line) # Check args validity param_card = self.check_check(args) @@ -9355,6 +9362,28 @@ def set2_apply_flavor_grouping(self, args, log=True): """ self.options['apply_flavor_grouping'] = banner_module.ConfigFile.format_variable(args[0], bool, 'apply_flavor_grouping') + def help_set2_color_basis(self): + logger.info("color_basis ",'$MG:color:GREEN') + logger.info(" > (default: auto) select the color basis used for processes") + logger.info(" whose color structure is purely adjoint (multi-gluon).") + logger.info(" > trace: the (n-1)! basis of traces of fundamental generators") + logger.info(" > ddm: the (n-2)! Del Duca-Dixon-Maltoni half-ladder basis") + logger.info(" (n-1 times fewer JAMPs, (n-1)^2 times smaller color matrix)") + logger.info(" > auto: ddm for the output formats which do not need a color") + logger.info(" flow decomposition (standalone), trace otherwise") + + def set2_color_basis(self, args, log=True): + """Set the color basis used for fully adjoint (multi-gluon) processes. + Example: set color_basis ddm + """ + args = ['color_basis'] + args + self.check_set(args) + value = args[1].lower() + if value not in self._valid_color_basis: + raise self.InvalidCmd('color_basis needs one of %s, got %s' % \ + (self._valid_color_basis, args[1])) + self.options['color_basis'] = value + # not documented options: @@ -9798,10 +9827,47 @@ def do_output(self, line): self._export_dir = None # Export a matrix element - def export(self, nojpeg = False, main_file_name = "", group_processes=True, + def set_color_basis_mode(self, *exporters): + """Set the color basis used for fully adjoint (multi-gluon) processes. + The (n-2)! Del Duca-Dixon-Maltoni basis can only be used by the output + formats which never need a color flow decomposition, so in 'auto' mode + every exporter involved must support it.""" + + mode = self.options.get('color_basis', 'auto') + exporters = [exporter for exporter in exporters if exporter] + if mode == 'auto': + use_ddm = bool(exporters) and \ + all(getattr(exporter, 'support_ddm_color_basis', False) + for exporter in exporters) + else: + use_ddm = (mode == 'ddm') + + # An exporter which has to write a color flow per event also needs the + # trace basis next to the DDM one + with_flow = any(getattr(exporter, 'ddm_needs_flow_basis', False) + for exporter in exporters) + + color_amp.set_ddm_basis(use_ddm, with_flow=with_flow) + if use_ddm: + logger.debug('Using the Del Duca-Dixon-Maltoni color basis for ' + 'fully adjoint processes (flow basis: %s)', with_flow) + + def export(self, nojpeg = False, main_file_name = "", group_processes=True, args=[]): """Export a generated amplitude to file.""" + self.set_color_basis_mode(self._curr_exporter, self._me_curr_exporter) + try: + return self._export(nojpeg, main_file_name, group_processes, args) + finally: + # the color basis is tied to this output, it must not leak to the + # next command + color_amp.set_ddm_basis(False) + + def _export(self, nojpeg = False, main_file_name = "", group_processes=True, + args=[]): + """Export a generated amplitude to file, with the color basis already + selected.""" # Define the helas call writer if hasattr(self._curr_exporter, 'helas_exporter') and self._curr_exporter.helas_exporter: diff --git a/madgraph/interface/master_interface.py b/madgraph/interface/master_interface.py index b18de52528..ae8649dee6 100755 --- a/madgraph/interface/master_interface.py +++ b/madgraph/interface/master_interface.py @@ -627,6 +627,9 @@ def help_set2_nlo_mixed_expansion(self, *args, **opts): def help_set2_output_dependencies(self, *args, **opts): return self.cmd.help_set2_output_dependencies(self, *args, **opts) + def help_set2_color_basis(self, *args, **opts): + return self.cmd.help_set2_color_basis(self, *args, **opts) + def help_set2_zerowidth_tchannel(self, *args, **opts): return self.cmd.help_set2_zerowidth_tchannel(self, *args, **opts) diff --git a/madgraph/iolibs/export_v4.py b/madgraph/iolibs/export_v4.py index 36bd3979e2..407ad8c62a 100755 --- a/madgraph/iolibs/export_v4.py +++ b/madgraph/iolibs/export_v4.py @@ -109,6 +109,13 @@ class VirtualExporter(object): exporter = 'v4' # language of the output 'v4' for Fortran output # 'cpp' for C++ output + support_ddm_color_basis = False + # True for the output formats which can use the (n-2)! Del Duca-Dixon- + # Maltoni basis for the color sum of multi-gluon processes. + ddm_needs_flow_basis = False + # True when the format also needs a color flow per event: the trace basis + # is then built next to the DDM one, and the trace JAMPs are obtained from + # the DDM ones through the Kleiss-Kuijf relations. default_vector_size = 0 @@ -1775,8 +1782,8 @@ def get_leshouche_lines(self, matrix_element, numproc): * (-1)**(1+l.get('state')) # Get the list of color flows color_flow_list = \ - matrix_element.get('color_basis').color_flow_decomposition(repr_dict, - ninitial) + matrix_element.get('color_basis').get_flow_basis().\ + color_flow_decomposition(repr_dict, ninitial) # And output them properly for cf_i, color_flow_dict in enumerate(color_flow_list): for i in [0, 1]: @@ -2382,6 +2389,56 @@ def get_den_factor_line(self, matrix_element): return "DATA IDEN/%2r/" % \ matrix_element.get_denominator_factor() + def set_color_flow_lines(self, matrix_element, replace_dict, ncolor): + """Fill in replace_dict everything the matrix element template needs to + know about the color flow basis, and return its size. + + For a fully adjoint (multi-gluon) process the color sum can be done on + the (n-2)! Del Duca-Dixon-Maltoni basis, but a color flow still has to + be picked among the (n-1)! trace structures. The trace JAMPs are then + not built from the amplitudes but obtained from the DDM ones through + the Kleiss-Kuijf relations, which is (n-1) times cheaper.""" + + color_basis = matrix_element.get('color_basis') + flow_basis = color_basis.get_flow_basis() if color_basis else None + + if flow_basis is None or flow_basis is color_basis: + replace_dict['ncolor_flow'] = ncolor + replace_dict['jampflow_decl'] = '' + replace_dict['jampflow_lines'] = '' + replace_dict['jamp_flow'] = 'JAMP' + return ncolor + + ncolor_flow = max(1, len(flow_basis)) + projection = color_basis.get_flow_projection() + + # The Kleiss-Kuijf map only acts on color, so it is the same for every + # split order + lines = [] + cmd_options = dict(self.cmd_options) + self.cmd_options['jamp_optim'] = False + try: + for iso in range(replace_dict['nAmpSplitOrders']): + flow_lines, nb_temp = self.get_JAMP_lines(projection, + JAMP_format="JAMPF(%%s,%d)" % (iso + 1), + AMP_format="JAMP(%%s,%d)" % (iso + 1)) + lines.extend(flow_lines) + finally: + self.cmd_options = cmd_options + + replace_dict['ncolor_flow'] = ncolor_flow + replace_dict['jampflow_decl'] = \ + ' COMPLEX*16 JAMPF(NCOLOR_FLOW,NAMPSO)' + replace_dict['jampflow_lines'] = '\n'.join(lines) + replace_dict['jamp_flow'] = 'JAMPF' + + logger.debug('Color sum on %d DDM structures, color flow on %d trace ' + 'structures (%d Kleiss-Kuijf terms)', + ncolor, ncolor_flow, + sum(len(row) for row in projection)) + + return ncolor_flow + def get_icolamp_lines(self, mapconfigs, matrix_element, num_matrix_element): """Return the ICOLAMP matrix, showing which JAMPs contribute to which configs (diagrams).""" @@ -2404,20 +2461,21 @@ def get_icolamp_lines(self, mapconfigs, matrix_element, num_matrix_element): # There is a color basis - create a list showing which JAMPs have # contributions to which configs - # Only want to include leading color flows, so find max_Nc - color_basis = matrix_element.get('color_basis') - + # Only want to include leading color flows, so find max_Nc. This is + # about color flows, so always the trace basis + color_basis = matrix_element.get('color_basis').get_flow_basis() + # We don't want to include the power of Nc's which come from the potential # loop color trace (i.e. in the case of a closed fermion loop for example) # so we subtract it here when computing max_Nc - max_Nc = max(sum([[(v[4]-v[5]) for v in val] for val in + max_Nc = max(sum([[(v[4]-v[5]) for v in val] for val in color_basis.values()],[])) # Crate dictionary between diagram number and JAMP number diag_jamp = {} for ijamp, col_basis_elem in \ - enumerate(sorted(matrix_element.get('color_basis').keys())): - for diag_tuple in matrix_element.get('color_basis')[col_basis_elem]: + enumerate(sorted(color_basis.keys())): + for diag_tuple in color_basis[col_basis_elem]: # Only use color flows with Nc == max_Nc. However, notice that # we don't want to include the Nc power coming from the loop # in this counting. @@ -5088,6 +5146,8 @@ class ProcessExporterFortranSA(ProcessExporterFortran): jamp_fold = True jamp_orbit = True default_vector_size = 0 + # standalone only squares the amplitude, it never writes color flows + support_ddm_color_basis = True # When True, emit per-call IAND(WF_FLAVOR_MASK/AMP_FLAVOR_MASK, # CURRENT_FLAV_BIT) guards in MATRIX so that wavefunctions and amplitudes # which contribute zero for the current input flavor are skipped at @@ -6365,6 +6425,8 @@ class ProcessExporterFortranMatchBox(ProcessExporterFortranSA): matrix_template = "matrix_standalone_matchbox.inc" + # matchbox needs the color flow information + support_ddm_color_basis = False @staticmethod def get_color_string_lines(matrix_element): @@ -8126,6 +8188,10 @@ def write_matrix_element_v4(self, writer, matrix_element, fortran_model, replace_dict['jamp_tmp_decl'] = '' if recipes else \ " COMPLEX*16 TMP_JAMP(%i)" % nb_temp + # The color sum can run on the (n-2)! DDM basis while the color flow + # probabilities keep using the (n-1)! trace one + ncolor = self.set_color_flow_lines(matrix_element, replace_dict, ncolor) + if self.beam_polarization == [True, True]: replace_dict['beam_polarization'] = """ DO JJ=1,nincoming @@ -9411,6 +9477,10 @@ class ProcessExporterFortranMEGroup(ProcessExporterFortranME): MadEvent subprocess group format.""" + # the color sum uses the DDM basis, the color flows the trace one + support_ddm_color_basis = True + ddm_needs_flow_basis = True + matrix_file = "matrix_madevent_group_v4.inc" grouped_mode = 'madevent' default_opt = {'clean': False, 'complex_mass':False, diff --git a/madgraph/iolibs/template_files/matrix_madevent_group_v4.inc b/madgraph/iolibs/template_files/matrix_madevent_group_v4.inc index bcce9e7522..cd421f7b0c 100644 --- a/madgraph/iolibs/template_files/matrix_madevent_group_v4.inc +++ b/madgraph/iolibs/template_files/matrix_madevent_group_v4.inc @@ -136,7 +136,7 @@ C ---------- AMP2(I)=0D0 ENDDO ENDIF - JAMP2(0)=%(ncolor)d + JAMP2(0)=%(ncolor_flow)d DO I=1,INT(JAMP2(0)) JAMP2(I)=0D0 ENDDO @@ -299,6 +299,8 @@ C INTEGER NWAVEFUNCS, NCOLOR, NCOLORFOLD PARAMETER (NWAVEFUNCS=%(nwavefuncs)d, NCOLOR=%(ncolor)d) PARAMETER (NCOLORFOLD=%(ncolorfold)d) + INTEGER NCOLOR_FLOW + PARAMETER (NCOLOR_FLOW=%(ncolor_flow)d) REAL*8 ZERO PARAMETER (ZERO=0D0) COMPLEX*16 IMAG1 @@ -330,6 +332,7 @@ C COMMON /%(proc_prefix)scolor_matrix%(proc_id)s/ CF,DENOM COMPLEX*16 AMP(NGRAPHS), JAMP(NCOLOR,NAMPSO) %(color_fold_decl)s +%(jampflow_decl)s type(aloha) W(NWAVEFUNCS) C Needed for v4 models COMPLEX*16 DUM0,DUM1 @@ -391,6 +394,7 @@ AMP(:) = (0d0,0d0) JAMP(:,:) = (0d0,0d0) %(jamp_lines)s +%(jampflow_lines)s if(init_mode)then DO I=1, NGRAPHS @@ -423,11 +427,11 @@ JAMP(:,:) = (0d0,0d0) %(amp2_lines)s endif - Do I = 1, NCOLOR + Do I = 1, NCOLOR_FLOW DO M = 1, NAMPSO DO N = 1, NAMPSO %(select_configs_if)s - Jamp2(i)=Jamp2(i)+DABS(DBLE(Jamp(i,m)*dconjg(Jamp(i,n)))) + Jamp2(i)=Jamp2(i)+DABS(DBLE(%(jamp_flow)s(i,m)*dconjg(%(jamp_flow)s(i,n)))) %(select_configs_endif)s enddo enddo diff --git a/madgraph/iolibs/template_files/matrix_madevent_group_v4_hel.inc b/madgraph/iolibs/template_files/matrix_madevent_group_v4_hel.inc index 6c6cd8b94a..922123457b 100644 --- a/madgraph/iolibs/template_files/matrix_madevent_group_v4_hel.inc +++ b/madgraph/iolibs/template_files/matrix_madevent_group_v4_hel.inc @@ -103,7 +103,7 @@ C ---------- AMP2(I)=0D0 ENDDO ENDIF - JAMP2(0)=%(ncolor)d + JAMP2(0)=%(ncolor_flow)d DO I=1,INT(JAMP2(0)) JAMP2(I)=0D0 ENDDO @@ -210,6 +210,8 @@ C INTEGER NWAVEFUNCS, NCOLOR, NCOLORFOLD PARAMETER (NWAVEFUNCS=${nwavefuncs}, NCOLOR=%(ncolor)d) PARAMETER (NCOLORFOLD=%(ncolorfold)d) + INTEGER NCOLOR_FLOW + PARAMETER (NCOLOR_FLOW=%(ncolor_flow)d) REAL*8 ZERO PARAMETER (ZERO=0D0) COMPLEX*16 IMAG1 @@ -244,6 +246,7 @@ C COMMON /%(proc_prefix)scolor_matrix%(proc_id)s/ CF,DENOM COMPLEX*16 AMP(NCOMB,NGRAPHS), JAMP(NCOLOR,NAMPSO) %(color_fold_decl)s +%(jampflow_decl)s type(aloha) W(NWAVEFUNCS) C Needed for v4 models COMPLEX*16 DUM0,DUM1 @@ -296,6 +299,7 @@ ${helas_calls} DO K = 1, NCOMB ${jamp_lines} %(color_fold_gather)s +%(jampflow_lines)s TS(K) = 0.D0 DO M = 1, NAMPSO CF_INDEX = 0 @@ -316,11 +320,11 @@ ${jamp_lines} if(sde_strat.eq.1) then ${amp2_lines} endif - Do I = 1, NCOLOR + Do I = 1, NCOLOR_FLOW DO M = 1, NAMPSO DO N = 1, NAMPSO %(select_configs_if)s - Jamp2(i)=Jamp2(i)+DABS(DBLE(Jamp(i,m)*dconjg(Jamp(i,n)))) + Jamp2(i)=Jamp2(i)+DABS(DBLE(%(jamp_flow)s(i,m)*dconjg(%(jamp_flow)s(i,n)))) %(select_configs_endif)s enddo enddo diff --git a/tests/unit_tests/core/test_color_algebra.py b/tests/unit_tests/core/test_color_algebra.py index bf65214c42..da87178e21 100755 --- a/tests/unit_tests/core/test_color_algebra.py +++ b/tests/unit_tests/core/test_color_algebra.py @@ -166,6 +166,11 @@ def test_f_object(self): self.assertEqual(my_f.simplify(), color.ColorFactor([col_str1, col_str2])) + # f is real, so complex conjugation must NOT reverse the indices + # (which would flip the sign of the totally antisymmetric f) + self.assertEqual(color.f(1, 2, 3).complex_conjugate(), + color.f(1, 2, 3)) + def test_d_object(self): """Test the d color object""" # T should have exactly 3 indices! diff --git a/tests/unit_tests/core/test_color_amp.py b/tests/unit_tests/core/test_color_amp.py index b507eb7987..63604254a0 100755 --- a/tests/unit_tests/core/test_color_amp.py +++ b/tests/unit_tests/core/test_color_amp.py @@ -17,6 +17,7 @@ color information for diagrams.""" from __future__ import absolute_import +import collections import copy import fractions @@ -997,3 +998,133 @@ def test_spanning_tree_rebuilds_color_matrix(self): self.assertEqual(node, representative[line]) self.assertEqual([rows[node][perm[j]] for j in range(len(keys))], rows[line]) + + +class DDMColorBasisTest(unittest.TestCase): + """Test the Del Duca-Dixon-Maltoni (n-2)! adjoint color basis""" + + mypartlist = base_objects.ParticleList() + myinterlist = base_objects.InteractionList() + mymodel = base_objects.Model() + + def setUp(self): + # same gluon + quark model as ColorSquareTest, built only once + if not len(self.mypartlist): + ColorSquareTest.setUp(self) + color_amp.set_ddm_basis(True) + + def tearDown(self): + color_amp.set_ddm_basis(False) + + def get_amplitude(self, ids): + """Amplitude for the process with the given pdg codes, the first two + being in the initial state.""" + + myleglist = base_objects.LegList() + for i, pdg in enumerate(ids): + myleglist.append(base_objects.Leg({'id': pdg, 'state': i > 1})) + + myamplitude = diagram_generation.Amplitude() + myamplitude.set('process', base_objects.Process({'legs': myleglist, + 'model': self.mymodel})) + myamplitude.generate_diagrams() + + return myamplitude + + def test_ddm_half_ladder(self): + """The half-ladder structures written down explicitly""" + + self.assertEqual(color_amp.ddm_half_ladder((2,), 1, 3), + (('f', (1, 2, 3)),)) + self.assertEqual(color_amp.ddm_half_ladder((3, 2), 1, 4), + (('f', (-1, 2, 4)), ('f', (1, 3, -1)))) + self.assertEqual(color_amp.ddm_half_ladder((2, 3, 4), 1, 5), + (('f', (-2, 4, 5)), ('f', (-1, 3, -2)), + ('f', (1, 2, -1)))) + + def test_ddm_basis_size_multi_gluons(self): + """The DDM basis of gg > n*g has (n+2-2)! elements instead of (n+1)!""" + + for n, size in enumerate([1, 2, 6, 24]): + amplitude = self.get_amplitude([21] * (n + 3)) + self.assertEqual(len(color_amp.ColorBasis(amplitude)), size) + + def test_ddm_reduction_matches_trace_basis(self): + """Expanding the DDM decomposition of every diagram back on the trace + basis must give the direct trace decomposition.""" + + def trace_expand(col_str): + res = collections.defaultdict(fractions.Fraction) + for cs in color.ColorFactor([col_str]).full_simplify(): + res[(cs.to_immutable(), cs.is_imaginary, cs.Nc_power)] += \ + cs.coeff + return dict((k, v) for k, v in res.items() if v) + + for nb_gluons in range(3, 6): + amplitude = self.get_amplitude([21] * nb_gluons) + col_basis = color_amp.ColorBasis() + color_dicts = col_basis.create_color_dict_list(amplitude) + ends = col_basis.get_ddm_ends() + self.assertEqual(ends, (1, nb_gluons)) + + nb_checked = 0 + for color_dict in color_dicts: + for col_str in color_dict.values(): + via_ddm = collections.defaultdict(fractions.Fraction) + for ddm_str in color_amp.reduce_to_ddm(col_str, *ends): + for k, v in trace_expand(ddm_str).items(): + via_ddm[k] += v + self.assertEqual(trace_expand(col_str), + dict((k, v) for k, v in via_ddm.items() if v)) + nb_checked += 1 + self.assertTrue(nb_checked > 0) + + def test_ddm_color_matrix_gg_gg(self): + """The 2x2 DDM color matrix of gg > gg, N^2(N^2-1) on the diagonal and + half of it off diagonal""" + + col_basis = color_amp.ColorBasis(self.get_amplitude([21] * 4)) + col_matrix = color_amp.ColorMatrix(col_basis, Nc=3) + + self.assertEqual(len(col_basis), 2) + for i, j, goal in [(0, 0, 72), (1, 1, 72), (0, 1, 36), (1, 0, 36)]: + self.assertEqual(col_matrix.col_matrix_fixed_Nc[(i, j)], + (fractions.Fraction(goal, 1), 0)) + + def test_ddm_color_matrix_is_positive_definite(self): + """A Gram matrix of real color tensors: the diagonal must be positive + and the matrix symmetric, for an even as well as an odd number of + gluons.""" + + for nb_gluons in range(3, 6): + col_basis = color_amp.ColorBasis(self.get_amplitude([21] * nb_gluons)) + col_matrix = color_amp.ColorMatrix(col_basis, Nc=3) + for i in range(len(col_basis)): + real, imag = col_matrix.col_matrix_fixed_Nc[(i, i)] + self.assertTrue(real > 0) + self.assertEqual(imag, 0) + for j in range(len(col_basis)): + self.assertEqual(col_matrix.col_matrix_fixed_Nc[(i, j)], + col_matrix.col_matrix_fixed_Nc[(j, i)]) + + def test_ddm_fall_back_on_trace_basis(self): + """Processes which are not fully adjoint keep the trace basis""" + + # u u~ > g g has T objects in its color structures + col_basis = color_amp.ColorBasis(self.get_amplitude([2, -2, 21, 21])) + self.assertEqual(col_basis._ddm_ends, None) + self.assertEqual(len(col_basis), 2) + + # and so has u u~ > u u~ + col_basis = color_amp.ColorBasis(self.get_amplitude([2, -2, 2, -2])) + self.assertEqual(col_basis._ddm_ends, None) + self.assertEqual(len(col_basis), 2) + + def test_ddm_no_color_flow_decomposition(self): + """The color flow decomposition is not defined in the DDM basis, and + must say so instead of returning something wrong.""" + + col_basis = color_amp.ColorBasis(self.get_amplitude([21] * 4)) + self.assertRaises(color_amp.ColorBasis.ColorBasisError, + col_basis.color_flow_decomposition, + {1: 8, 2: 8, 3: 8, 4: 8}, 2) From 7487ad5baf2777028a03201baf807879d719dbac Mon Sep 17 00:00:00 2001 From: Olivier Mattelaer Date: Fri, 7 Aug 2026 18:51:45 +0200 Subject: [PATCH 33/42] give standalone the same color flow work as madevent Standalone is the mode used to time the matrix element, so it has to carry what madevent carries. It now builds the trace basis next to the DDM one and emits GET_JAMPF, the Kleiss-Kuijf rebuild of the trace JAMPs, together with the JAMP2 accumulation that madevent does for its color flow probabilities. The batched BLAS color sum takes its own branch in SMATRIX and never enters MATRIX, so the flow JAMPs are computed there as well; without that the BLAS timings were measuring a matrix element madevent could not use. Co-Authored-By: Claude Opus 5 --- madgraph/iolibs/export_v4.py | 118 ++++++++++++++++-- .../template_files/matrix_standalone_v4.inc | 4 + 2 files changed, 109 insertions(+), 13 deletions(-) diff --git a/madgraph/iolibs/export_v4.py b/madgraph/iolibs/export_v4.py index 407ad8c62a..d091d02b59 100755 --- a/madgraph/iolibs/export_v4.py +++ b/madgraph/iolibs/export_v4.py @@ -2389,6 +2389,78 @@ def get_den_factor_line(self, matrix_element): return "DATA IDEN/%2r/" % \ matrix_element.get_denominator_factor() + def get_flow_jamp_lines(self, projection, JAMP_format, AMP_format): + """The Kleiss-Kuijf definitions of the trace JAMPs in terms of the DDM + ones. The common subexpression pass is skipped: the map has only a + handful of terms per line and its temporaries would collide with the + ones of the JAMP definitions proper.""" + + cmd_options = dict(self.cmd_options) + self.cmd_options['jamp_optim'] = False + try: + lines, nb_temp = self.get_JAMP_lines(projection, + JAMP_format=JAMP_format, + AMP_format=AMP_format) + finally: + self.cmd_options = cmd_options + + return lines + + def set_color_flow_lines_sa(self, matrix_element, replace_dict, ncolor): + """Same as set_color_flow_lines, for the standalone template: the JAMPs + are not split per amplitude order there, and the flow JAMPs get their + own routine so that they can be timed on their own.""" + + prefix = replace_dict['proc_prefix'] + color_basis = matrix_element.get('color_basis') + flow_basis = color_basis.get_flow_basis() if color_basis else None + + if flow_basis is None or flow_basis is color_basis: + replace_dict['ncolor_flow'] = ncolor + replace_dict['jampflow_decl'] = '' + replace_dict['jampflow_call'] = '' + replace_dict['jampflow_routine'] = '' + return ncolor + + ncolor_flow = max(1, len(flow_basis)) + projection = color_basis.get_flow_projection() + lines = self.get_flow_jamp_lines(projection, JAMP_format="JAMPF(%s)", + AMP_format="JAMP(%s)") + + replace_dict['ncolor_flow'] = ncolor_flow + replace_dict['jampflow_decl'] = "\n".join([ + " INTEGER NCOLOR_FLOW", + " PARAMETER (NCOLOR_FLOW=%d)" % ncolor_flow, + " COMPLEX*16 JAMPF(NCOLOR_FLOW)", + " DOUBLE PRECISION %sJAMP2(NCOLOR_FLOW)" % prefix, + " COMMON /%sJAMP2_COMMON/ %sJAMP2" % (prefix, prefix)]) + # accumulated exactly like madevent does, so that the work is real + replace_dict['jampflow_call'] = "\n".join([ + " CALL %sGET_JAMPF(JAMP,JAMPF)" % prefix, + " DO I = 1, NCOLOR_FLOW", + " %sJAMP2(I) = %sJAMP2(I)" % (prefix, prefix), + " $ + DABS(DBLE(JAMPF(I)*DCONJG(JAMPF(I))))", + " ENDDO"]) + replace_dict['jampflow_routine'] = "\n".join([ + " SUBROUTINE %sGET_JAMPF(JAMP,JAMPF)" % prefix, + "CF2PY INTENT(OUT) :: JAMPF", + "CF2PY INTENT(IN) :: JAMP", + " IMPLICIT NONE", + " INTEGER NCOLOR, NCOLOR_FLOW", + " PARAMETER (NCOLOR=%d)" % ncolor, + " PARAMETER (NCOLOR_FLOW=%d)" % ncolor_flow, + " COMPLEX*16 IMAG1", + " PARAMETER (IMAG1=(0D0,1D0))", + " COMPLEX*16 JAMP(NCOLOR), JAMPF(NCOLOR_FLOW)"] + lines + + [" END"]) + + logger.debug('Color sum on %d DDM structures, color flow on %d trace ' + 'structures (%d Kleiss-Kuijf terms)', + ncolor, ncolor_flow, + sum(len(row) for row in projection)) + + return ncolor_flow + def set_color_flow_lines(self, matrix_element, replace_dict, ncolor): """Fill in replace_dict everything the matrix element template needs to know about the color flow basis, and return its size. @@ -2415,16 +2487,10 @@ def set_color_flow_lines(self, matrix_element, replace_dict, ncolor): # The Kleiss-Kuijf map only acts on color, so it is the same for every # split order lines = [] - cmd_options = dict(self.cmd_options) - self.cmd_options['jamp_optim'] = False - try: - for iso in range(replace_dict['nAmpSplitOrders']): - flow_lines, nb_temp = self.get_JAMP_lines(projection, - JAMP_format="JAMPF(%%s,%d)" % (iso + 1), - AMP_format="JAMP(%%s,%d)" % (iso + 1)) - lines.extend(flow_lines) - finally: - self.cmd_options = cmd_options + for iso in range(replace_dict['nAmpSplitOrders']): + lines.extend(self.get_flow_jamp_lines(projection, + JAMP_format="JAMPF(%%s,%d)" % (iso + 1), + AMP_format="JAMP(%%s,%d)" % (iso + 1))) replace_dict['ncolor_flow'] = ncolor_flow replace_dict['jampflow_decl'] = \ @@ -5146,8 +5212,11 @@ class ProcessExporterFortranSA(ProcessExporterFortran): jamp_fold = True jamp_orbit = True default_vector_size = 0 - # standalone only squares the amplitude, it never writes color flows + # standalone only squares the amplitude, so it can use the DDM basis. It + # still carries the Kleiss-Kuijf reconstruction of the trace JAMPs, because + # this is the mode used to time the matrix element and madevent pays it. support_ddm_color_basis = True + ddm_needs_flow_basis = True # When True, emit per-call IAND(WF_FLAVOR_MASK/AMP_FLAVOR_MASK, # CURRENT_FLAV_BIT) guards in MATRIX so that wavefunctions and amplitudes # which contribute zero for the current input flavor are skipped at @@ -6143,6 +6212,11 @@ def write_matrix_element_v4(self, writer, matrix_element, fortran_model, replace_dict['jamp_lines'] = '\n'.join(jamp_lines) + # The color flow JAMPs, rebuilt from the ones entering the color sum. + # Standalone does not need a color flow, but it is the mode used to + # time the matrix element, so it carries the same work as madevent. + self.set_color_flow_lines_sa(matrix_element, replace_dict, ncolor) + # The definitions written as one recipe per orbit are held in one # array together with the amplitudes, so that the loop running them # reads its two operands from the same place. @@ -6165,6 +6239,23 @@ def write_matrix_element_v4(self, writer, matrix_element, fortran_model, prefix = replace_dict['proc_prefix'] reps = ([line + 1 for line in folding['representatives']] if folding else list(range(1, ncolor + 1))) + # The batch branch does not go through MATRIX, so it has to carry the + # color flow JAMPs itself, exactly like the per-helicity path does + nflow = replace_dict.get('ncolor_flow', ncolor) + flow_decl, flow_lines = [], [] + if replace_dict.get('jampflow_routine'): + flow_decl = [" COMPLEX*16 JAMPFB(%d)" % nflow, + " DOUBLE PRECISION %sJAMP2(%d)" % (prefix, nflow), + " COMMON /%sJAMP2_COMMON/ %sJAMP2" % (prefix, + prefix)] + flow_lines = [ + " CALL %sGET_JAMPF(JAMPB,JAMPFB)" % prefix, + " DO IBH = 1, %d" % nflow, + " %sJAMP2(IBH) = %sJAMP2(IBH)" % (prefix, prefix), + " $ + DABS(DBLE(JAMPFB(IBH)*DCONJG(JAMPFB(IBH" + "))))", + " ENDDO"] + if self.blas_wanted(nfold): replace_dict['blas_guard_open'] = "(" replace_dict['blas_guard'] = ") .AND. .NOT.BLASDONE" @@ -6177,7 +6268,7 @@ def write_matrix_element_v4(self, writer, matrix_element, fortran_model, if recipes else 0), ncolor), " DOUBLE PRECISION, ALLOCATABLE, SAVE :: JRB(:,:)", " DOUBLE PRECISION, ALLOCATABLE, SAVE :: JIB(:,:)", - " INTEGER COLREPB(%d)" % nfold] + + " INTEGER COLREPB(%d)" % nfold] + flow_decl + self.get_int_data_lines("COLREPB", reps, var='IBH')) replace_dict['blas_branch'] = "\n".join([ " BLASDONE = .FALSE.", @@ -6193,7 +6284,8 @@ def write_matrix_element_v4(self, writer, matrix_element, fortran_model, " NBHEL = NBHEL + 1", " CALL %sGET_AMP(P,NHEL(1,IHEL),JC(1),FLAV_IDX,AMPB)" % prefix, - " CALL %sGET_JAMP(AMPB,JAMPB)" % prefix, + " CALL %sGET_JAMP(AMPB,JAMPB)" % prefix] + + flow_lines + [ " DO IBH = 1, %d" % nfold, " JRB(IBH,NBHEL) = DBLE(JAMPB(COLREPB(IBH)))", " JIB(IBH,NBHEL) = DIMAG(JAMPB(COLREPB(IBH)))", diff --git a/madgraph/iolibs/template_files/matrix_standalone_v4.inc b/madgraph/iolibs/template_files/matrix_standalone_v4.inc index f93d9db8e0..5c512ee9ac 100644 --- a/madgraph/iolibs/template_files/matrix_standalone_v4.inc +++ b/madgraph/iolibs/template_files/matrix_standalone_v4.inc @@ -232,6 +232,7 @@ C common /%(proc_prefix)scolor_matrix/ %(proc_prefix)sCF,%(proc_prefix)sDENOM COMPLEX*16 AMP(%(namp_dim)s), JAMP(NCOLOR) %(jamp_tmp_decl)s +%(jampflow_decl)s type(aloha) W(NWAVEFUNCS) COMPLEX*16 DUM0,DUM1 DATA DUM0, DUM1/(0d0, 0d0), (1d0, 0d0)/ @@ -251,6 +252,7 @@ C ---------- c WRITE (*,*) ' -> AMP = ', AMP call %(proc_prefix)sGET_JAMP(AMP,JAMP) c WRITE (*,*) ' -> JAMP = ', JAMP +%(jampflow_call)s call %(proc_prefix)sGET_MATRIX(JAMP,%(proc_prefix)sMATRIX) c write (*,*) " -> col.ave. |M|^2 for HEL=[", NHEL ,"] = ", %(proc_prefix)sMATRIX @@ -354,6 +356,8 @@ C %(jamp_lines)s END +%(jampflow_routine)s + SUBROUTINE %(proc_prefix)sGET_MATRIX(JAMP,MATRIX) C %(process_lines)s From 6f559c36e64894c9e3b74f50b3cb36bd764dda7e Mon Sep 17 00:00:00 2001 From: Olivier Mattelaer Date: Fri, 7 Aug 2026 22:04:28 +0200 Subject: [PATCH 34/42] give madmatrix the DDM color basis and the Kleiss-Kuijf color flows The C++ backend now does for a multi-gluon process what the fortran ones do: the color sum runs on the (n-2)! Del Duca-Dixon-Maltoni structures, while the color flow is still picked among the (n-1)! trace ones, whose jamps are rebuilt from the DDM ones through the Kleiss-Kuijf relations rather than from the amplitudes. The two counts were the same number before, so ncolor did for both. They are now separate: ncolor stays the color sum, ncolor_flow is what a color flow is picked among, and every buffer holding jamp2 -- the accumulation in calculate_jamps, the jamp2_sv arrays, the colAllJamp2s super-buffer, the selection walking icolamp -- moves to ncolor_flow. Without a DDM basis ncolor_flow is ncolor and the generated code is what it was. g g > g g g g: ncolor 120 -> 24, CPPProcess.cc 512 -> 269 kB, |M|^2 unchanged to the last digit (1.5929925846563478e-04 against 1.5929925846563475e-04). Co-Authored-By: Claude Opus 5 --- madgraph/iolibs/export_cpp.py | 13 +++++ madgraph/iolibs/export_mg7.py | 7 ++- .../madmatrix/MatrixElementKernels.cc | 2 +- .../madmatrix/process_class.inc | 3 ++ .../process_function_definitions.inc | 5 +- .../madmatrix/process_matrix.inc | 9 ++-- .../madmatrix/process_sigmaKin_function.inc | 18 +++---- madmatrix/model_handling.py | 49 +++++++++++++++++++ madmatrix/output.py | 8 +++ 9 files changed, 96 insertions(+), 18 deletions(-) diff --git a/madgraph/iolibs/export_cpp.py b/madgraph/iolibs/export_cpp.py index 354815b369..babeae4ca1 100755 --- a/madgraph/iolibs/export_cpp.py +++ b/madgraph/iolibs/export_cpp.py @@ -1665,6 +1665,10 @@ def get_matrix_single_process(self, i, matrix_element, color_amplitudes, replace_dict['jamp_lines'] = self.get_jamp_lines(color_amplitudes) + # The color sum may run on a smaller basis than the one the color flow + # is picked among (see the madmatrix override) + self.set_color_flow_lines_cpp(matrix_element, replace_dict) + replace_dict['amp2_lines'] = self.get_amp2_lines(matrix_element) #specific exporter hack @@ -1817,6 +1821,15 @@ def coeff(cls, ff_number, frac, is_imaginary, Nc_power, Nc_value=3): + def set_color_flow_lines_cpp(self, matrix_element, replace_dict): + """Tell the process template that the color sum and the color flow use + the same basis. Overridden by the backends which can put the color sum + on a smaller one.""" + + replace_dict['ncolor_flow'] = replace_dict['ncolor'] + replace_dict['jampflow_lines'] = '' + replace_dict['jamp_flow'] = 'jamp_sv' + def get_jamp_lines(self, color_amplitudes): """Return the jamp = sum(fermionfactor * amp[i]) lines""" diff --git a/madgraph/iolibs/export_mg7.py b/madgraph/iolibs/export_mg7.py index a6cc4f0518..a9b5e458c3 100644 --- a/madgraph/iolibs/export_mg7.py +++ b/madgraph/iolibs/export_mg7.py @@ -205,8 +205,11 @@ def get_subprocess_info(self, proc_dir, lib_me_path): repr_dict[leg.get("number")] = self.model.get_particle( leg.get("id") ).get_color() * (-1) ** (1 + leg.get("state")) - # Get the list of color flows - color_flow_dicts = self.color_basis.color_flow_decomposition(repr_dict, n_initial) + # Get the list of color flows. This is about color flows, so + # always the trace basis, even when the color sum runs on the DDM + # one. + color_flow_dicts = self.color_basis.get_flow_basis().\ + color_flow_decomposition(repr_dict, n_initial) # And output them properly color_flows = [ [[color_flow_dict[leg.get("number")][i] for i in [0, 1]] for leg in legs] diff --git a/madgraph/iolibs/template_files/madmatrix/MatrixElementKernels.cc b/madgraph/iolibs/template_files/madmatrix/MatrixElementKernels.cc index 872e4795e3..5cc43344be 100644 --- a/madgraph/iolibs/template_files/madmatrix/MatrixElementKernels.cc +++ b/madgraph/iolibs/template_files/madmatrix/MatrixElementKernels.cc @@ -310,7 +310,7 @@ namespace mg5amcGpu , m_pHelJamps() , m_pHelNumerators() , m_pHelDenominators() - , m_colJamp2s( CPPProcess::ncolor * this->nevt() ) + , m_colJamp2s( CPPProcess::ncolor_flow * this->nevt() ) #ifdef MGONGPU_CHANNELID_DEBUG , m_hstChannelIds( this->nevt() ) #endif diff --git a/madgraph/iolibs/template_files/madmatrix/process_class.inc b/madgraph/iolibs/template_files/madmatrix/process_class.inc index 59a6d07333..87c6242721 100644 --- a/madgraph/iolibs/template_files/madmatrix/process_class.inc +++ b/madgraph/iolibs/template_files/madmatrix/process_class.inc @@ -56,6 +56,9 @@ static constexpr int ncomb = %(nbhel)d; // #helicity combinations: e.g. 16 for e+ e- -> mu+ mu- (2**4 = fermion spin up/down ** npar) static constexpr int ndiagrams = %(ndiagrams)d; // #Feynman diagrams: e.g. 3 for e+ e- -> mu+ mu- 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 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 082c373aa2..eaef9990c8 100644 --- a/madgraph/iolibs/template_files/madmatrix/process_function_definitions.inc +++ b/madgraph/iolibs/template_files/madmatrix/process_function_definitions.inc @@ -83,6 +83,7 @@ namespace mg5amcCpu constexpr int npar = CPPProcess::npar; // #particles in total (external = initial + final): e.g. 4 for e+ e- -> mu+ mu- 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 // [NB: I am currently unable to get the right value of nwf in CPPProcess.h - will hardcode it in CPPProcess.cc instead (#644)] @@ -770,7 +771,7 @@ namespace mg5amcCpu const fptype* allrndcol, // input: random numbers[nevt] for color selection const fptype* allrnddiagram, // input: random numbers[nevt] for diagram selection const unsigned int* allChannelIds, // input: multichannel channelIds[nevt] (1 to #diagrams); nullptr to disable SDE enhancement (fix #899/#911) - const fptype_sv* allJamp2s, // input: jamp2[ncolor][nevt] for color choice (nullptr if disabled) + const fptype_sv* allJamp2s, // input: jamp2[ncolor_flow][nevt] for color choice (nullptr if disabled) const fptype* allNumerators, // input: all numerators const fptype* allDenominators, // input: all denominators const int nevt ) // input: #events (for cuda: nevt == ndim == gpublocks*gputhreads) @@ -811,7 +812,7 @@ namespace mg5amcCpu assert( channelId <= mgOnGpu::nchannels ); // SANITY CHECK #919 #910 } // Determine the jamp2 for this event (TEMPORARY? could do this with a dedicated memory accessor instead...) - fptype_sv jamp2_sv[ncolor] = { 0 }; + fptype_sv jamp2_sv[ncolor_flow] = { 0 }; assert( allJamp2s != nullptr ); // sanity check using J2_ACCESS = DeviceAccessJamp2; for( int icolC = 0; icolC < ncolor; icolC++ ) diff --git a/madgraph/iolibs/template_files/madmatrix/process_matrix.inc b/madgraph/iolibs/template_files/madmatrix/process_matrix.inc index dff402fa3c..438365dfb2 100644 --- a/madgraph/iolibs/template_files/madmatrix/process_matrix.inc +++ b/madgraph/iolibs/template_files/madmatrix/process_matrix.inc @@ -8,23 +8,24 @@ ! Integrated with the MadGraph7 project in Feb 2026. !========================================================================== +%(jampflow_lines)s // *** COLOR CHOICE BELOW *** // Store the leading color flows for choice of color #ifndef MGONGPUCPP_GPUIMPL if( jamp2_sv ) // disable color choice if nullptr { - for( int icol = 0; icol < ncolor; icol++ ) - jamp2_sv[ncolor * iParity + icol] += cxabs2( jamp_sv[icol] ); // may underflow #831 + for( int icol = 0; icol < ncolor_flow; icol++ ) + jamp2_sv[ncolor_flow * iParity + icol] += cxabs2( %(jamp_flow)s[icol] ); // may underflow #831 } #else /* clang-format off */ assert( iParity == 0 ); // sanity check for J2_ACCESS using J2_ACCESS = DeviceAccessJamp2; if( colAllJamp2s ) // disable color choice if nullptr { - for( int icol = 0; icol < ncolor; icol++ ) + 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_sv[icol] ) ); + atomicAdd( &J2_ACCESS::kernelAccessIcol( colAllJamp2s, icol ), cxabs2( %(jamp_flow)s[icol] ) ); } #endif /* clang-format on */ diff --git a/madgraph/iolibs/template_files/madmatrix/process_sigmaKin_function.inc b/madgraph/iolibs/template_files/madmatrix/process_sigmaKin_function.inc index 227301a6ea..dc39791d85 100644 --- a/madgraph/iolibs/template_files/madmatrix/process_sigmaKin_function.inc +++ b/madgraph/iolibs/template_files/madmatrix/process_sigmaKin_function.inc @@ -16,7 +16,7 @@ const int nevt = gpublocks * gputhreads; gpuMemset( allMEs, 0, nevt * sizeof( fptype ) ); gpuMemset( ghelAllJamps, 0, cNGoodHel * ncolor * mgOnGpu::nx2 * nevt * sizeof( fptype ) ); - gpuMemset( colAllJamp2s, 0, ncolor * nevt * sizeof( fptype ) ); + gpuMemset( colAllJamp2s, 0, ncolor_flow * nevt * sizeof( fptype ) ); gpuMemset( ghelAllNumerators, 0, cNGoodHel * processConfig::ndiagrams * nevt * sizeof( fptype ) ); gpuMemset( ghelAllDenominators, 0, cNGoodHel * nevt * sizeof( fptype ) ); gpuMemset( ghelAllMEs, 0, cNGoodHel * nevt * sizeof( fptype ) ); @@ -110,8 +110,8 @@ const int ievt00 = ipagV2 * neppV; // loop on one SIMD page (neppV events) at a time #endif // Running sum of partial amplitudes squared for event by event color selection (#402) - // (jamp2[nParity][ncolor][neppV] for the SIMD vector - or the two SIMD vectors - of events processed in calculate_jamps) - fptype_sv jamp2_sv[nParity * ncolor] = {}; + // (jamp2[nParity][ncolor_flow][neppV] for the SIMD vector - or the two SIMD vectors - of events processed in calculate_jamps) + fptype_sv jamp2_sv[nParity * ncolor_flow] = {}; fptype_sv MEs_ighel[ncomb] = {}; // sum of MEs for all good helicities up to ighel (for the first - and/or only - neppV page) #if defined MGONGPU_CPPSIMD and defined MGONGPU_FPTYPE_DOUBLE and defined MGONGPU_FPTYPE2_FLOAT fptype_sv MEs_ighel2[ncomb] = {}; // sum of MEs for all good helicities up to ighel (for the second neppV page) @@ -235,9 +235,9 @@ printf( "INTERNAL ERROR! Cannot choose an event-by-event random color for channelId=%%d (invalid SDE iconfig=%%d\n > nconfig=%%d)", channelId, iconfig, mgOnGpu::nconfigSDE ); assert( iconfig <= (int)mgOnGpu::nconfigSDE ); // SANITY CHECK #917 } - fptype targetamp[ncolor] = { 0 }; + fptype targetamp[ncolor_flow] = { 0 }; // NB (see #877): explicitly use 'icolC' rather than 'icol' to indicate that icolC uses C indexing in [0, N_colors-1] - for( int icolC = 0; icolC < ncolor; icolC++ ) + for( int icolC = 0; icolC < ncolor_flow; icolC++ ) { if( icolC == 0 ) targetamp[icolC] = 0; @@ -245,17 +245,17 @@ targetamp[icolC] = targetamp[icolC - 1]; #ifdef MGONGPU_CPPSIMD if( mgOnGpu::icolamp[iconfig - 1][icolC] ) targetamp[icolC] += - jamp2_sv[icolC + ncolor * ( ieppV / neppV )][ieppV %% neppV]; + jamp2_sv[icolC + ncolor_flow * ( ieppV / neppV )][ieppV %% neppV]; #else if( mgOnGpu::icolamp[iconfig - 1][icolC] ) targetamp[icolC] += - jamp2_sv[icolC + ncolor * ( ieppV / neppV )]; + jamp2_sv[icolC + ncolor_flow * ( ieppV / neppV )]; #endif } const int ievt = ievt00 + ieppV; //printf( "sigmaKin: ievt=%%4d rndcol=%%f\n", ievt, allrndcol[ievt] ); - for( int icolC = 0; icolC < ncolor; icolC++ ) + for( int icolC = 0; icolC < ncolor_flow; icolC++ ) { - if( allrndcol[ievt] < ( targetamp[icolC] / targetamp[ncolor - 1] ) ) + if( allrndcol[ievt] < ( targetamp[icolC] / targetamp[ncolor_flow - 1] ) ) { allselcol[ievt] = icolC + 1; // NB Fortran [1,ncolor], cudacpp [0,ncolor-1] //printf( "sigmaKin: ievt=%%d icol=%%d\n", ievt, icolC+1 ); diff --git a/madmatrix/model_handling.py b/madmatrix/model_handling.py index b4d1bac787..71e6c62064 100644 --- a/madmatrix/model_handling.py +++ b/madmatrix/model_handling.py @@ -1503,6 +1503,9 @@ def __init__(self, *args, **kwargs): def get_process_class_definitions(self, write=True): replace_dict = super().get_process_class_definitions(write=False) replace_dict['process_lines'] = replace_dict['process_lines'].replace('\n','\n ') + # ncolor_flow sits next to ncolor in the class, so it has to be known + # here as well as in get_process_function_definitions + self.set_color_flow_lines_cpp(self.matrix_elements[0], replace_dict) ###misc.sprint( replace_dict['nwavefuncs'] ) # NB: this (from export_cpp) is the WRONG value of nwf, e.g. 6 for gg_tt (#644) ###misc.sprint( self.matrix_elements[0].get_number_of_wavefunctions() ) # NB: this is a different WRONG value of nwf, e.g. 7 for gg_tt (#644) ###replace_dict['nwavefunc'] = self.matrix_elements[0].get_number_of_wavefunctions() # how do I get HERE the right value of nwf, e.g. 5 for gg_tt? @@ -1729,6 +1732,9 @@ def get_process_function_definitions(self, write=True): replace_dict['all_flavors'] = replace_dict['all_flavors'].replace('flavors', 'tFlavors') color_amplitudes = [me.get_color_amplitudes() for me in self.matrix_elements] # as in OneProcessExporterCPP.get_process_function_definitions replace_dict['ncolor'] = len(color_amplitudes[0]) + # The color sum can run on the (n-2)! DDM basis while the color flow + # probabilities keep using the (n-1)! trace one + self.set_color_flow_lines_cpp(self.matrix_elements[0], replace_dict) # broken_symmetry_factor function: use the shared decay-aware symmetry # data (same as the Fortran / standalone_cpp exporters) instead of the # old simple PID-count version, so identical-particle and decay-chain @@ -2128,6 +2134,49 @@ def write_process_cc_file(self, writer): return replace_dict # AV - replace the export_cpp.OneProcessExporterCPP method (fix fptype and improve formatting) + def set_color_flow_lines_cpp(self, matrix_element, replace_dict): + """Fill in replace_dict everything the process template needs to know + about the color flow basis. + + For a fully adjoint (multi-gluon) process the color sum can be done on + the (n-2)! Del Duca-Dixon-Maltoni basis, but a color flow still has to + be picked among the (n-1)! trace structures. The trace jamps are then + not built from the amplitudes but obtained from the DDM ones through + the Kleiss-Kuijf relations, which is (n-1) times cheaper.""" + + color_basis = matrix_element.get('color_basis') + flow_basis = color_basis.get_flow_basis() if color_basis else None + + 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' + return + + projection = color_basis.get_flow_projection() + 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) + for coefficient, number in coeff_list) + lines.append(' jampf_sv[%d] = %s;' % (i, terms if terms + else 'cxzero_sv()')) + + replace_dict['ncolor_flow'] = max(1, len(flow_basis)) + replace_dict['jampflow_lines'] = '\n'.join(lines) + replace_dict['jamp_flow'] = 'jampf_sv' + + logger.debug('Color sum on %d DDM structures, color flow on %d trace ' + 'structures (%d Kleiss-Kuijf terms)', + replace_dict['ncolor'], replace_dict['ncolor_flow'], + sum(len(row) for row in projection)) + def get_color_matrix_lines(self, matrix_element): """Return the color matrix definition lines for this matrix element. Split rows in chunks of size n.""" import madgraph.core.color_algebra as color diff --git a/madmatrix/output.py b/madmatrix/output.py index 1820b10e96..9dc7113547 100644 --- a/madmatrix/output.py +++ b/madmatrix/output.py @@ -54,6 +54,14 @@ 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 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 + # trace jamps are rebuilt from the DDM ones through the Kleiss-Kuijf + # relations (see set_color_flow_lines_cpp in model_handling.py). + support_ddm_color_basis = True + ddm_needs_flow_basis = True + # Below are the class variable that are defined in export_cpp.ProcessExporterGPU # AV - keep defaults from export_cpp.ProcessExporterGPU # Decide which type of merging is used [madevent/madweight] From ed00d1a8bb4fbbf08c12dab845a8f969f0ac1816 Mon Sep 17 00:00:00 2001 From: Olivier Mattelaer Date: Fri, 7 Aug 2026 22:54:20 +0200 Subject: [PATCH 35/42] fold the madmatrix color sum onto one flow per reversal pair Reversing a color flow gives the same flow back up to one overall sign, so half of them carry nothing of their own and |M|^2 can be summed over one flow per pair against a color matrix folded onto them. The fortran output has done this for a while; the C++/GPU one summed the full ncolor x ncolor triangle. The reflection and the folded matrix are not reimplemented: the four methods the fortran exporters were carrying move to a ColorReflectionFolding mixin, which madmatrix inherits. The only thing the two backends disagree on is when a folding is worth taking, now a jamp_fold_worthwhile hook: fortran can rebuild a sign +1 matrix at run time and so refuses to write a large sign -1 one, while madmatrix always writes the matrix out and folded it is a quarter of the size. What is generated is the folded matrix alone, over ncolorfold flows, plus the colorFoldRep table saying which flow of each pair is kept. color_sum_cpu and color_sum_kernel gather through that table and sum over ncolorfold. jamp2 is untouched: color selection needs every flow separately, and it is accumulated from the full local jamp before the gather. The BLAS color sum is NOT folded - folding it means compacting allJamps, which the jamps are not written in - but it does not get a second matrix either: it multiplies the folded matrix spread back over the ncolor flows, with the dropped rows and columns left at zero, so it cannot drift away from what the kernel computes. Compacting the jamps is left to do. Measured on the color sum alone (FPTYPE=d, one call, arm64): g g > g g g 24 -> 12 flows, sign -1: 95 ns -> 28 ns (3.4x) g g > g g g g 120 -> 60 flows, sign +1: 3.6 us -> 0.71 us (5.1x) Validated against unfolded builds over 512 phase space points: max relative difference 6.5e-16 (5 ulp) for g g > g g g and 1.2e-15 (8 ulp) for g g > g g g g -- a re-grouped sum cannot be bit-identical, and a wrong sign is not a 1e-16 effect. Both agree with the fortran ./check to the same accuracy, and the folded matrices match the exact rational fold of the unfolded ones on every entry. u u~ > u u~ g, which does not fold, stays bit-identical. Co-Authored-By: Claude Opus 5 --- madgraph/iolibs/export_v4.py | 325 ++++++++++-------- .../template_files/madmatrix/color_sum.cc | 164 +++++---- madmatrix/model_handling.py | 104 +++++- 3 files changed, 357 insertions(+), 236 deletions(-) diff --git a/madgraph/iolibs/export_v4.py b/madgraph/iolibs/export_v4.py index b99c7137f4..48add083fa 100755 --- a/madgraph/iolibs/export_v4.py +++ b/madgraph/iolibs/export_v4.py @@ -220,10 +220,179 @@ def export_helas(self, HELAS_PATH): raise Exception("V4 model not supported by this type of exporter. Please use UFO model") return +#=============================================================================== +# ColorReflectionFolding +#=============================================================================== +class ColorReflectionFolding(object): + """Reversing every color basis element maps the basis onto itself, and for + a pure gluon process the two flows of a pair only differ by one overall + sign. Half the color flows then carry nothing of their own and |M|^2 can be + summed over one flow per pair, against a color matrix folded onto them. + + Shared by the fortran exporters and by the madmatrix (C++/GPU) one, which + only differ in when a folding is worth taking (jamp_fold_worthwhile).""" + + # sum |M|^2 over one color flow per reversal pair instead of over every one + # Folding the color matrix onto one line per reversal pair only works + # where the template sums over NCOLORFOLD. get_color_data_lines is shared + # by every fortran exporter, so this stays off unless the template agrees. + jamp_fold = False + + # Above this many entries the folded color matrix is not written out but + # rebuilt at run time, which only the sign +1 case can do (see + # jamp_fold_worthwhile). + color_fold_max_written = 300000 + + @staticmethod + def jamp_color_rows(matrix_element): + """The color coefficient of every amplitude, one dictionary per color + basis line. Same numbers get_JAMP_lines works from.""" + + rows = [] + powers = {} + for coeff_list in matrix_element.get_color_amplitudes(): + row = {} + for coefficient, amp in coeff_list: + if not coefficient: + continue + try: + power = powers[coefficient[3]] + except KeyError: + power = fractions.Fraction(3) ** coefficient[3] + powers[coefficient[3]] = power + value = (1j if coefficient[2] else 1) * coefficient[0] * \ + coefficient[1] * power + row[amp] = row.get(amp, 0) + value + rows.append(dict((amp, complex(v)) for amp, v in row.items() if v)) + return rows + + def get_jamp_reflection(self, matrix_element): + """Reversing every color basis element maps the basis onto itself, and + for a pure gluon process the color coefficients of a line and of its + reverse differ by one overall sign, so half the color flows carry no + information of their own: + + JAMP[reverse(i)] = sign * JAMP[i] + + Return (reverse, sign) or None. The relation is read off the color + coefficients themselves rather than assumed, so a process where it does + not hold -- a quark line, where reversing does not commute with the + fermion flow -- simply gets None.""" + + if not isinstance(matrix_element, helas_objects.HelasMatrixElement): + return None + color_basis = matrix_element.get('color_basis') + if not color_basis or len(color_basis) < 2: + return None + keys = sorted(color_basis.keys()) + position = dict((key, i) for i, key in enumerate(keys)) + + reverse = [] + for key in keys: + other = color_amp.reverse_immutable(key) + if other is None or other not in position: + return None + reverse.append(position[other]) + if any(reverse[reverse[i]] != i for i in range(len(keys))): + return None + + columns = self.jamp_color_rows(matrix_element) + + sign = None + for i in range(len(keys)): + here, there = columns[i], columns[reverse[i]] + if set(here) != set(there): + return None + for amp, value in here.items(): + ratio = there[amp] / value + if ratio not in (1, -1): + return None + if sign is None: + sign = int(ratio.real) + elif sign != int(ratio.real): + return None + if sign is None: + return None + return reverse, sign + + @staticmethod + def jamp_reflection_representatives(reverse): + """One line per pair, and for every line the pair it belongs to.""" + + representatives = [i for i in range(len(reverse)) if i <= reverse[i]] + slot = {} + for index, line in enumerate(representatives): + slot[line] = index + slot[reverse[line]] = index + return representatives, slot + + def jamp_fold_worthwhile(self, sign, nb_pairs): + """Whether a folding is taken once it has been found. + + With sign +1 every line of a pair enters with the same weight, so the + permutations leaving the color basis invariant carry over to the pairs + unchanged and the folded matrix can still be rebuilt at run time from + one line per orbit. With sign -1 a permutation may send a line onto its + own partner, which flips the weight, and the rebuilt form would need a + sign of its own; there the folded matrix is written out instead, which + is only affordable while it stays small.""" + + return sign > 0 or \ + nb_pairs * (nb_pairs + 1) // 2 <= self.color_fold_max_written + + def get_jamp_folding(self, matrix_element): + """Whether to sum |M|^2 over one line per reversal pair, and the + (reverse, sign, representatives, slot) that goes with it.""" + + if not self.jamp_fold: + return None + found = self.get_jamp_reflection(matrix_element) + if not found: + return None + reverse, sign = found + representatives, slot = self.jamp_reflection_representatives(reverse) + if not self.jamp_fold_worthwhile(sign, len(representatives)): + return None + return {'reverse': reverse, 'sign': sign, + 'representatives': representatives, 'slot': slot} + + def jamp_folded_color_matrix(self, matrix_element, reverse, sign): + """The color matrix over one line per reversal pair. Summing |M|^2 over + the pairs instead of over every line gives the same number, since the + two lines of a pair only differ by the overall sign: + + C'[a][b] = sum over the two lines of a and the two of b, each + weighted by its sign relative to the line kept + + Returns (denominator, rows) with rows[a][b] integer, a and b indexing + the representatives.""" + + color_matrix = matrix_element.get('color_matrix') + representatives, _slot = self.jamp_reflection_representatives(reverse) + denominator = max(color_matrix.get_line_denominators()) + full = [color_matrix.get_line_numerators(i, denominator) + for i in range(len(reverse))] + + def pair(a): + return [(a, 1)] if reverse[a] == a else [(a, 1), (reverse[a], sign)] + + rows = [] + for a in representatives: + row = [] + for b in representatives: + total = 0 + for i, ci in pair(a): + for j, cj in pair(b): + total += ci * cj * full[i][j] + assert int(total) == total + row.append(int(total)) + rows.append(row) + return denominator, rows + #=============================================================================== # ProcessExporterFortran #=============================================================================== -class ProcessExporterFortran(VirtualExporter): +class ProcessExporterFortran(ColorReflectionFolding, VirtualExporter): """Class to take care of exporting a set of matrix elements to Fortran (v4) format.""" @@ -236,11 +405,9 @@ class ProcessExporterFortran(VirtualExporter): jamp_optim = False # how many times the JAMP optimisation called itself, for the record myjamp_count = 0 - # sum |M|^2 over one color flow per reversal pair instead of over every one - # Folding the color matrix onto one line per reversal pair only works - # where the template sums over NCOLORFOLD. get_color_data_lines is shared - # by every fortran exporter, so this stays off unless the template agrees. - jamp_fold = False + # jamp_fold (sum |M|^2 over one color flow per reversal pair) comes from + # ColorReflectionFolding and stays off unless the template sums over + # NCOLORFOLD: get_color_data_lines is shared by every fortran exporter. jamp_integer_walk = True # BLAS-3 for the color sum: all helicities at once as one right hand side. # None means take it when the library is there and the process is big @@ -3283,83 +3450,6 @@ def optimise_jamp_best(self, all_element, symmetry): # does, leaves the matrix invariant at every step, and the definitions can # be written as one recipe per orbit. - @staticmethod - def jamp_color_rows(matrix_element): - """The color coefficient of every amplitude, one dictionary per color - basis line. Same numbers get_JAMP_lines works from.""" - - rows = [] - powers = {} - for coeff_list in matrix_element.get_color_amplitudes(): - row = {} - for coefficient, amp in coeff_list: - if not coefficient: - continue - try: - power = powers[coefficient[3]] - except KeyError: - power = fractions.Fraction(3) ** coefficient[3] - powers[coefficient[3]] = power - value = (1j if coefficient[2] else 1) * coefficient[0] * \ - coefficient[1] * power - row[amp] = row.get(amp, 0) + value - rows.append(dict((amp, complex(v)) for amp, v in row.items() if v)) - return rows - - def get_jamp_reflection(self, matrix_element): - """Reversing every color basis element maps the basis onto itself, and - for a pure gluon process the color coefficients of a line and of its - reverse differ by one overall sign, so half the color flows carry no - information of their own: - - JAMP[reverse(i)] = sign * JAMP[i] - - Return (reverse, sign) or None. The relation is read off the color - coefficients themselves rather than assumed, so a process where it does - not hold -- a quark line, where reversing does not commute with the - fermion flow -- simply gets None.""" - - if not isinstance(matrix_element, helas_objects.HelasMatrixElement): - return None - color_basis = matrix_element.get('color_basis') - if not color_basis or len(color_basis) < 2: - return None - keys = sorted(color_basis.keys()) - position = dict((key, i) for i, key in enumerate(keys)) - - reverse = [] - for key in keys: - other = color_amp.reverse_immutable(key) - if other is None or other not in position: - return None - reverse.append(position[other]) - if any(reverse[reverse[i]] != i for i in range(len(keys))): - return None - - columns = self.jamp_color_rows(matrix_element) - - sign = None - for i in range(len(keys)): - here, there = columns[i], columns[reverse[i]] - if set(here) != set(there): - return None - for amp, value in here.items(): - ratio = there[amp] / value - if ratio not in (1, -1): - return None - if sign is None: - sign = int(ratio.real) - elif sign != int(ratio.real): - return None - if sign is None: - return None - return reverse, sign - - # Above this many entries the folded color matrix is not written out but - # rebuilt at run time, which only the sign +1 case can do (see - # get_jamp_folding). - color_fold_max_written = 300000 - _blas_available = None @classmethod @@ -3439,31 +3529,6 @@ def jamp_global_phase(all_element): return None return phase - def get_jamp_folding(self, matrix_element): - """Whether to sum |M|^2 over one line per reversal pair, and the - (reverse, sign, representatives, slot) that goes with it. - - With sign +1 every line of a pair enters with the same weight, so the - permutations leaving the color basis invariant carry over to the pairs - unchanged and the folded matrix can still be rebuilt at run time from - one line per orbit. With sign -1 a permutation may send a line onto its - own partner, which flips the weight, and the rebuilt form would need a - sign of its own; there the folded matrix is written out instead, which - is only affordable while it stays small.""" - - if not self.jamp_fold: - return None - found = self.get_jamp_reflection(matrix_element) - if not found: - return None - reverse, sign = found - representatives, slot = self.jamp_reflection_representatives(reverse) - nb = len(representatives) - if sign < 0 and nb * (nb + 1) // 2 > self.color_fold_max_written: - return None - return {'reverse': reverse, 'sign': sign, - 'representatives': representatives, 'slot': slot} - @staticmethod def get_blas_routine(prefix, nfold, ncomb): """The color sum for every helicity at once. DSYMM is real, so the @@ -3545,50 +3610,6 @@ def get_color_fold_ampso(self, folding, ncolor): 'color_fold_gather': " JFOLD(:,:) = JAMP(COLREP(:),:)", 'color_fold_array': 'JFOLD'} - def jamp_folded_color_matrix(self, matrix_element, reverse, sign): - """The color matrix over one line per reversal pair. Summing |M|^2 over - the pairs instead of over every line gives the same number, since the - two lines of a pair only differ by the overall sign: - - C'[a][b] = sum over the two lines of a and the two of b, each - weighted by its sign relative to the line kept - - Returns (denominator, rows) with rows[a][b] integer, a and b indexing - the representatives.""" - - color_matrix = matrix_element.get('color_matrix') - representatives, _slot = self.jamp_reflection_representatives(reverse) - denominator = max(color_matrix.get_line_denominators()) - full = [color_matrix.get_line_numerators(i, denominator) - for i in range(len(reverse))] - - def pair(a): - return [(a, 1)] if reverse[a] == a else [(a, 1), (reverse[a], sign)] - - rows = [] - for a in representatives: - row = [] - for b in representatives: - total = 0 - for i, ci in pair(a): - for j, cj in pair(b): - total += ci * cj * full[i][j] - assert int(total) == total - row.append(int(total)) - rows.append(row) - return denominator, rows - - @staticmethod - def jamp_reflection_representatives(reverse): - """One line per pair, and for every line the pair it belongs to.""" - - representatives = [i for i in range(len(reverse)) if i <= reverse[i]] - slot = {} - for index, line in enumerate(representatives): - slot[line] = index - slot[reverse[line]] = index - return representatives, slot - @staticmethod def jamp_column_form(column): """Canonical form of one column of the JAMP matrix up to a global sign, diff --git a/madgraph/iolibs/template_files/madmatrix/color_sum.cc b/madgraph/iolibs/template_files/madmatrix/color_sum.cc index 30c6799932..5340b9f5c0 100644 --- a/madgraph/iolibs/template_files/madmatrix/color_sum.cc +++ b/madgraph/iolibs/template_files/madmatrix/color_sum.cc @@ -24,21 +24,42 @@ namespace mg5amcCpu %(color_matrix_lines)s #ifdef MGONGPUCPP_GPUIMPL - // The normalized color matrix (divide each column by denom) + // The normalized folded color matrix (divide each column by denom) template struct NormalizedColorMatrix { constexpr __host__ __device__ NormalizedColorMatrix() : value() { - for( int icol = 0; icol < ncolor; icol++ ) - for( int jcol = 0; jcol < ncolor; jcol++ ) - value[icol * ncolor + jcol] = colorMatrix[icol][jcol] / colorDenom[icol]; + for( int ifold = 0; ifold < ncolorfold; ifold++ ) + for( int jfold = 0; jfold < ncolorfold; jfold++ ) + value[ifold * ncolorfold + jfold] = colorMatrix[ifold][jfold] / colorDenom[ifold]; + } + T value[ncolorfold * ncolorfold]; + }; + // The fptype2 version is the default used by kernels (supporting mixed floating point mode) + static __device__ fptype2 s_pNormalizedColorMatrixFold2[ncolorfold * ncolorfold]; +#ifndef MGONGPU_HAS_NO_BLAS + // The same matrix spread back over the ncolor unfolded color flows, which is what BLAS + // multiplies: the rows and columns of the flows which are not kept are left at zero, so + // the dropped flows contribute nothing and the product is the folded sum written out in + // full. The BLAS color sum is therefore NOT folded - it does the same ncolor x ncolor + // work it did before - but it takes its numbers from the one folded matrix which is + // written out, so it cannot drift away from what the kernel computes. + template + struct UnfoldedNormalizedColorMatrix + { + constexpr __host__ __device__ UnfoldedNormalizedColorMatrix() + : value() + { + for( int ifold = 0; ifold < ncolorfold; ifold++ ) + for( int jfold = 0; jfold < ncolorfold; jfold++ ) + value[hostColorFoldRep[ifold] * ncolor + hostColorFoldRep[jfold]] = colorMatrix[ifold][jfold] / colorDenom[ifold]; } T value[ncolor * ncolor]; }; - // The fptype2 version is the default used by kernels (supporting mixed floating point mode also in blas) static __device__ fptype2 s_pNormalizedColorMatrix2[ncolor * ncolor]; +#endif #endif //-------------------------------------------------------------------------- @@ -51,7 +72,11 @@ namespace mg5amcCpu { first = false; constexpr NormalizedColorMatrix normalizedColorMatrix2; - gpuMemcpyToSymbol( s_pNormalizedColorMatrix2, normalizedColorMatrix2.value, ncolor * ncolor * sizeof( fptype2 ) ); + gpuMemcpyToSymbol( s_pNormalizedColorMatrixFold2, normalizedColorMatrix2.value, ncolorfold * ncolorfold * sizeof( fptype2 ) ); +#ifndef MGONGPU_HAS_NO_BLAS + constexpr UnfoldedNormalizedColorMatrix unfoldedNormalizedColorMatrix2; + gpuMemcpyToSymbol( s_pNormalizedColorMatrix2, unfoldedNormalizedColorMatrix2.value, ncolor * ncolor * sizeof( fptype2 ) ); +#endif } } #endif @@ -71,64 +96,62 @@ namespace mg5amcCpu __host__ __device__ constexpr TriangularNormalizedColorMatrix() : value() { - for( int icol = 0; icol < ncolor; icol++ ) + for( int ifold = 0; ifold < ncolorfold; ifold++ ) { // Diagonal terms - value[icol][icol] = colorMatrix[icol][icol] / colorDenom[icol]; + value[ifold][ifold] = colorMatrix[ifold][ifold] / colorDenom[ifold]; // Off-diagonal terms - for( int jcol = icol + 1; jcol < ncolor; jcol++ ) - value[icol][jcol] = 2 * colorMatrix[icol][jcol] / colorDenom[icol]; + for( int jfold = ifold + 1; jfold < ncolorfold; jfold++ ) + value[ifold][jfold] = 2 * colorMatrix[ifold][jfold] / colorDenom[ifold]; } } - fptype2 value[ncolor][ncolor]; + fptype2 value[ncolorfold][ncolorfold]; }; static constexpr auto cf2 = TriangularNormalizedColorMatrix(); // Use the property that M is a real matrix (see #475): // we can rewrite the quadratic form (A-iB)(M)(A+iB) as AMA - iBMA + iBMA + BMB = AMA + BMB // In addition, on C++ use the property that M is symmetric (see #475), - // and also use constexpr to compute "2*" and "/colorDenom[icol]" once and for all at compile time: + // and also use constexpr to compute "2*" and "/colorDenom[ifold]" once and for all at compile time: // we gain (not a factor 2...) in speed here as we only loop over the up diagonal part of the matrix. // Strangely, CUDA is slower instead, so keep the old implementation for the moment. fptype_sv deltaMEs = { 0 }; #if defined MGONGPU_CPPSIMD and defined MGONGPU_FPTYPE_DOUBLE and defined MGONGPU_FPTYPE2_FLOAT fptype_sv deltaMEs_next = { 0 }; - // Mixed mode: merge two neppV vectors into one neppV2 vector - fptype2_sv jampR_sv[ncolor]; - fptype2_sv jampI_sv[ncolor]; - for( int icol = 0; icol < ncolor; icol++ ) - { - jampR_sv[icol] = fpvmerge( cxreal( allJamp_sv[icol] ), cxreal( allJamp_sv[ncolor + icol] ) ); - jampI_sv[icol] = fpvmerge( cximag( allJamp_sv[icol] ), cximag( allJamp_sv[ncolor + icol] ) ); - } -#else - const cxtype_sv* jamp_sv = allJamp_sv; #endif - // Loop over icol - for( int icol = 0; icol < ncolor; icol++ ) + // Gather the color flows the sum runs over: one per reversal pair when the color basis + // folds (ncolorfold < ncolor), every flow otherwise (ncolorfold == ncolor, identity map). + // NB in mixed mode the two neppV vectors of allJamp_sv, at icol and at ncolor+icol, are + // two halves of the event page and not two colors: it is the color index inside each of + // them which is gathered, and the two are merged into one neppV2 vector. + fptype2_sv jampR_sv[ncolorfold]; + fptype2_sv jampI_sv[ncolorfold]; + for( int ifold = 0; ifold < ncolorfold; ifold++ ) { - // Diagonal terms + const int icol = colorFoldRep[ifold]; #if defined MGONGPU_CPPSIMD and defined MGONGPU_FPTYPE_DOUBLE and defined MGONGPU_FPTYPE2_FLOAT - fptype2_sv& jampRi_sv = jampR_sv[icol]; - fptype2_sv& jampIi_sv = jampI_sv[icol]; + jampR_sv[ifold] = fpvmerge( cxreal( allJamp_sv[icol] ), cxreal( allJamp_sv[ncolor + icol] ) ); + jampI_sv[ifold] = fpvmerge( cximag( allJamp_sv[icol] ), cximag( allJamp_sv[ncolor + icol] ) ); #else - fptype2_sv jampRi_sv = (fptype2_sv)( cxreal( jamp_sv[icol] ) ); - fptype2_sv jampIi_sv = (fptype2_sv)( cximag( jamp_sv[icol] ) ); + jampR_sv[ifold] = (fptype2_sv)( cxreal( allJamp_sv[icol] ) ); + jampI_sv[ifold] = (fptype2_sv)( cximag( allJamp_sv[icol] ) ); #endif - fptype2_sv ztempR_sv = cf2.value[icol][icol] * jampRi_sv; - fptype2_sv ztempI_sv = cf2.value[icol][icol] * jampIi_sv; - // Loop over jcol - for( int jcol = icol + 1; jcol < ncolor; jcol++ ) + } + // Loop over ifold + for( int ifold = 0; ifold < ncolorfold; ifold++ ) + { + // Diagonal terms + fptype2_sv& jampRi_sv = jampR_sv[ifold]; + fptype2_sv& jampIi_sv = jampI_sv[ifold]; + fptype2_sv ztempR_sv = cf2.value[ifold][ifold] * jampRi_sv; + fptype2_sv ztempI_sv = cf2.value[ifold][ifold] * jampIi_sv; + // Loop over jfold + for( int jfold = ifold + 1; jfold < ncolorfold; jfold++ ) { // Off-diagonal terms -#if defined MGONGPU_CPPSIMD and defined MGONGPU_FPTYPE_DOUBLE and defined MGONGPU_FPTYPE2_FLOAT - fptype2_sv& jampRj_sv = jampR_sv[jcol]; - fptype2_sv& jampIj_sv = jampI_sv[jcol]; -#else - fptype2_sv jampRj_sv = (fptype2_sv)( cxreal( jamp_sv[jcol] ) ); - fptype2_sv jampIj_sv = (fptype2_sv)( cximag( jamp_sv[jcol] ) ); -#endif - ztempR_sv += cf2.value[icol][jcol] * jampRj_sv; - ztempI_sv += cf2.value[icol][jcol] * jampIj_sv; + fptype2_sv& jampRj_sv = jampR_sv[jfold]; + fptype2_sv& jampIj_sv = jampI_sv[jfold]; + ztempR_sv += cf2.value[ifold][jfold] * jampRj_sv; + ztempI_sv += cf2.value[ifold][jfold] * jampIj_sv; } fptype2_sv deltaMEs2 = ( jampRi_sv * ztempR_sv + jampIi_sv * ztempI_sv ); // may underflow #831 #if defined MGONGPU_CPPSIMD and defined MGONGPU_FPTYPE_DOUBLE and defined MGONGPU_FPTYPE2_FLOAT @@ -167,40 +190,42 @@ namespace mg5amcCpu allJamps = allJamps + ighel * nevtIfAllHelicities; // Jamps for one specific helicity ighel } using J_ACCESS = DeviceAccessJamp; - fptype jampR[ncolor]; - fptype jampI[ncolor]; - for( int icol = 0; icol < ncolor; icol++ ) + // Gather the color flows the sum runs over: one per reversal pair when the color basis + // folds (ncolorfold < ncolor), every flow otherwise (ncolorfold == ncolor, identity map) + fptype jampR[ncolorfold]; + fptype jampI[ncolorfold]; + for( int ifold = 0; ifold < ncolorfold; ifold++ ) { constexpr int ihel0 = 0; // the input buffer allJamps already points to a specific helicity - cxtype jamp = J_ACCESS::kernelAccessIcolIhelNhelConst( allJamps, icol, ihel0, nGoodHel ); - jampR[icol] = jamp.real(); - jampI[icol] = jamp.imag(); + cxtype jamp = J_ACCESS::kernelAccessIcolIhelNhelConst( allJamps, colorFoldRep[ifold], ihel0, nGoodHel ); + jampR[ifold] = jamp.real(); + jampI[ifold] = jamp.imag(); } - // Loop over icol + // Loop over ifold fptype deltaMEs = { 0 }; - for( int icol = 0; icol < ncolor; icol++ ) + for( int ifold = 0; ifold < ncolorfold; ifold++ ) { fptype2 ztempR = { 0 }; fptype2 ztempI = { 0 }; - fptype2 jampRi = jampR[icol]; - fptype2 jampIi = jampI[icol]; - // OLD IMPLEMENTATION (ihel3: symmetric square matrix) - Loop over all jcol - //for( int jcol = 0; jcol < ncolor; jcol++ ) + fptype2 jampRi = jampR[ifold]; + fptype2 jampIi = jampI[ifold]; + // OLD IMPLEMENTATION (ihel3: symmetric square matrix) - Loop over all jfold + //for( int jfold = 0; jfold < ncolorfold; jfold++ ) //{ - // fptype2 jampRj = jampR[jcol]; - // fptype2 jampIj = jampI[jcol]; - // ztempR += s_pNormalizedColorMatrix2[icol * ncolor + jcol] * jampRj; // use fptype2 version of color matrix - // ztempI += s_pNormalizedColorMatrix2[icol * ncolor + jcol] * jampIj; // use fptype2 version of color matrix + // fptype2 jampRj = jampR[jfold]; + // fptype2 jampIj = jampI[jfold]; + // ztempR += s_pNormalizedColorMatrixFold2[ifold * ncolorfold + jfold] * jampRj; // use fptype2 version of color matrix + // ztempI += s_pNormalizedColorMatrixFold2[ifold * ncolorfold + jfold] * jampIj; // use fptype2 version of color matrix //} - // NEW IMPLEMENTATION #475 (ihel3p1: triangular lower diagonal matrix) - Loop over jcol < icol - ztempR += s_pNormalizedColorMatrix2[icol * ncolor + icol] * jampRi; // use fptype2 version of color matrix - ztempI += s_pNormalizedColorMatrix2[icol * ncolor + icol] * jampIi; // use fptype2 version of color matrix - for( int jcol = 0; jcol < icol; jcol++ ) + // NEW IMPLEMENTATION #475 (ihel3p1: triangular lower diagonal matrix) - Loop over jfold < ifold + ztempR += s_pNormalizedColorMatrixFold2[ifold * ncolorfold + ifold] * jampRi; // use fptype2 version of color matrix + ztempI += s_pNormalizedColorMatrixFold2[ifold * ncolorfold + ifold] * jampIi; // use fptype2 version of color matrix + for( int jfold = 0; jfold < ifold; jfold++ ) { - fptype2 jampRj = jampR[jcol]; - fptype2 jampIj = jampI[jcol]; - ztempR += 2 * s_pNormalizedColorMatrix2[icol * ncolor + jcol] * jampRj; // use fptype2 version of color matrix - ztempI += 2 * s_pNormalizedColorMatrix2[icol * ncolor + jcol] * jampIj; // use fptype2 version of color matrix + fptype2 jampRj = jampR[jfold]; + fptype2 jampIj = jampI[jfold]; + ztempR += 2 * s_pNormalizedColorMatrixFold2[ifold * ncolorfold + jfold] * jampRj; // use fptype2 version of color matrix + ztempI += 2 * s_pNormalizedColorMatrixFold2[ifold * ncolorfold + jfold] * jampIj; // use fptype2 version of color matrix } deltaMEs += ztempR * jampRi; deltaMEs += ztempI * jampIi; @@ -272,6 +297,13 @@ namespace mg5amcCpu { const int nevt = gpublocks * gputhreads; + // NB: unlike color_sum_cpu and color_sum_kernel, the BLAS color sum is NOT folded onto one + // color flow per reversal pair. Folding it would mean compacting allJamps from ncolor down + // to ncolorfold, which the jamps are not written in, so it would take a gather kernel and a + // second buffer. Instead the matrix it multiplies is the folded one spread back over the + // ncolor flows (see UnfoldedNormalizedColorMatrix): same numbers, same ncolor x ncolor work + // as before. Compacting the jamps would make this a factor 4 cheaper and is left to do. + // Get the address associated with the normalized color matrix in device memory static fptype2* devNormColMat = nullptr; if( !devNormColMat ) gpuGetSymbolAddress( (void**)&devNormColMat, s_pNormalizedColorMatrix2 ); diff --git a/madmatrix/model_handling.py b/madmatrix/model_handling.py index b4d1bac787..ae0b249fa2 100644 --- a/madmatrix/model_handling.py +++ b/madmatrix/model_handling.py @@ -1462,7 +1462,8 @@ def get_mg5_info_lines(self): # (NB: enable this via ProcessExporterMadMatrix.oneprocessclass in output.py) # (NB: use this directly also in MadMatrixUFOModelConverter.read_template_file) # (NB: use this directly also in MadMatrixGPUFOHelasCallWriter.super_get_matrix_element_calls) -class OneProcessExporterMadMatrix(export_mg7.OneProcessExporterMG7): +class OneProcessExporterMadMatrix(export_v4.ColorReflectionFolding, + export_mg7.OneProcessExporterMG7): # Class structure information # - object # - OneProcessExporterCPP(object) [in madgraph/iolibs/export_cpp.py] @@ -1470,6 +1471,9 @@ class OneProcessExporterMadMatrix(export_mg7.OneProcessExporterMG7): # - OneProcessExporterMadMatrix(OneProcessExporterCPP) # This class + # Sum |M|^2 over one color flow per reversal pair (color_sum.cc) + jamp_fold = True + # AV - change defaults from export_cpp.OneProcessExporterCPP cc_ext = 'cc' # create CPPProcess.cc process_dir = '.' @@ -2127,33 +2131,97 @@ def write_process_cc_file(self, writer): else: return replace_dict + # The folded color matrix is always written out, never rebuilt at run time + # as the fortran output can do, and folded it is a quarter of the size of + # the matrix which would be written otherwise. So take every folding found. + def jamp_fold_worthwhile(self, sign, nb_pairs): + return True + # AV - replace the export_cpp.OneProcessExporterCPP method (fix fptype and improve formatting) def get_color_matrix_lines(self, matrix_element): - """Return the color matrix definition lines for this matrix element. Split rows in chunks of size n.""" - import madgraph.core.color_algebra as color + """Return the color matrix definition lines for this matrix element. Split rows in chunks of size n. + + |M|^2 is summed over one color flow per reversal pair when the basis + allows it (see ColorReflectionFolding), so what is written out is the + color matrix folded onto those flows, together with the list of the + flows kept. Without a folding every flow is its own representative and + the matrix is the plain one.""" if not matrix_element.get('color_matrix'): - return '\n'.join([' static constexpr fptype2 colorDenom[1] = {1.};', 'static const fptype2 cf[1][1] = {1.};']) + return '\n'.join([ + self.get_color_fold_lines(None, 1), + ' static constexpr fptype2 colorDenom[1] = {1.};', + ' static constexpr fptype2 colorMatrix[1][1] = {1.};']) else: - color_denominators = matrix_element.get('color_matrix').\ - get_line_denominators() - denom_string = ' static constexpr fptype2 colorDenom[ncolor] = { %s }; // 1-D array[%i]' \ - % ( ', '.join(['%i' % denom for denom in color_denominators]), len(color_denominators) ) - matrix_strings = [] - for index, denominator in enumerate(color_denominators): - # Then write the numerators for the matrix elements - num_list = matrix_element.get('color_matrix').get_line_numerators(index, denominator) - matrix_strings.append('{ %s }' % ', '.join(['%d' % i for i in num_list])) - matrix_string = ' static constexpr fptype2 colorMatrix[ncolor][ncolor] = ' + folding = self.get_jamp_folding(matrix_element) + if folding: + denominator, rows = self.jamp_folded_color_matrix( + matrix_element, folding['reverse'], folding['sign']) + color_denominators = [denominator] * len(rows) + num_lists = rows + ncolor = len(folding['reverse']) + else: + color_denominators = matrix_element.get('color_matrix').\ + get_line_denominators() + num_lists = [matrix_element.get('color_matrix'). + get_line_numerators(index, denominator) + for index, denominator + in enumerate(color_denominators)] + ncolor = len(color_denominators) + nfold = len(color_denominators) + denom_string = ' static constexpr fptype2 colorDenom[ncolorfold] = { %s }; // 1-D array[%i]' \ + % ( ', '.join(['%i' % denom for denom in color_denominators]), nfold ) + matrix_strings = ['{ %s }' % ', '.join(['%d' % i for i in num_list]) + for num_list in num_lists] + matrix_string = ' static constexpr fptype2 colorMatrix[ncolorfold][ncolorfold] = ' if len( matrix_strings ) > 1: matrix_string += '{\n ' + ',\n '.join(matrix_strings) + ' };' else: matrix_string += '{ ' + matrix_strings[0] + ' };' - matrix_string += ' // 2-D array[%i][%i]' % ( len(color_denominators), len(color_denominators) ) - denom_comment = '\n // The color denominators (initialize all array elements, with ncolor=%i)\n // [NB do keep \'static\' for these constexpr arrays, see issue #283]\n' % len(color_denominators) - matrix_comment = '\n // The color matrix (initialize all array elements, with ncolor=%i)\n // [NB do keep \'static\' for these constexpr arrays, see issue #283]\n' % len(color_denominators) + matrix_string += ' // 2-D array[%i][%i]' % ( nfold, nfold ) + denom_comment = '\n // The color denominators (initialize all array elements, with ncolorfold=%i)\n // [NB do keep \'static\' for these constexpr arrays, see issue #283]\n' % nfold + matrix_comment = '\n // The color matrix (initialize all array elements, with ncolorfold=%i)\n // [NB do keep \'static\' for these constexpr arrays, see issue #283]\n' % nfold denom_string = denom_comment + denom_string matrix_string = matrix_comment + matrix_string - return '\n'.join([denom_string, matrix_string]) + return '\n'.join([self.get_color_fold_lines(folding, ncolor), + denom_string, matrix_string]) + + @staticmethod + def get_color_fold_lines(folding, ncolor): + """The number of color flows the sum runs over and which flow it keeps + out of every reversal pair. Without a folding this is every flow.""" + + if folding: + representatives = folding['representatives'] + comment = ( + '\n // Reversing a color flow gives the same flow back up to an overall sign\n' + ' // (JAMP[reverse(i)] = %+i * JAMP[i] here), so only one flow of each reversal\n' + ' // pair carries anything of its own: |M|^2 is summed over those, against the\n' + ' // color matrix folded onto them (see ColorReflectionFolding in export_v4.py).\n' + % folding['sign']) + else: + representatives = list(range(ncolor)) + comment = ( + '\n // Reversal does not map this color basis onto itself up to one overall\n' + ' // sign, so every color flow enters the sum on its own.\n') + chunks = [', '.join('%i' % line for line in representatives[start:start + 20]) + for start in range(0, len(representatives), 20)] + values = '{\n ' + ',\n '.join(chunks) + ' }' + # colorFoldRep is indexed at run time inside the GPU kernel, so it has to + # live in device memory, and a host copy is needed next to it: same split + # as channel2iconfig/hostChannel2iconfig in coloramps.h + return comment + \ + ' constexpr int ncolorfold = %i; // the number of color flows |M|^2 is summed over\n' % len(representatives) + \ + ' // Which color flow of each reversal pair is kept (C indexing, in [0, ncolor-1])\n' + \ + ' // (NB: this array is created on the host in C++ code and on the device in GPU code)\n' + \ + ' __device__ constexpr int colorFoldRep[ncolorfold] = %s; // 1-D array[%i]\n' \ + % (values, len(representatives)) + \ + '#ifdef MGONGPUCPP_GPUIMPL\n' + \ + ' // Host copy of the colorFoldRep array (needed to fold the color matrix at compile time)\n' + \ + ' constexpr int hostColorFoldRep[ncolorfold] = %s; // 1-D array[%i]\n' \ + % (values, len(representatives)) + \ + '#else\n' + \ + ' constexpr const int* hostColorFoldRep = colorFoldRep;\n' + \ + '#endif' # AV - replace the export_cpp.OneProcessExporterCPP method (improve formatting) def get_initProc_lines(self, matrix_element, color_amplitudes): From 4ca220cba70bae7c2b9d66028c7449a0bc6cb740 Mon Sep 17 00:00:00 2001 From: Olivier Mattelaer Date: Fri, 7 Aug 2026 22:55:50 +0200 Subject: [PATCH 36/42] give the madmatrix color flows the sub-expressions the fortran ones have The C++/cudacpp backend wrote one 'jamp_sv[i] += c*amp_sv[0]' per (color flow, amplitude) pair, which for g g > g g g g is 8160 statements and half a megabyte of CPPProcess.cc. The fortran side has had a common sub-expression pass over the same coefficient matrix for a long time (optimise_jamp), but the two shared no code. Everything about that pass which is not printing now lives in madgraph/iolibs/jamp_optimiser.py, mixed into both exporters. The fortran output is unchanged, byte for byte. The C++ emitter cannot print the definitions the way fortran does: fortran keeps AMP(NGRAPHS) to read the amplitudes back from, while here every amplitude passes through the single slot amp_sv[0] and is gone by the next diagram. So they are accumulated instead -- each amplitude is added into the definition that uses it while it is still there, and the definitions built from other definitions are emitted at the amplitude they become ready after. That keeps the temporaries down to one array next to jamp_sv, where storing every amplitude would have cost 510 more cxtype_sv for g g > g g g g. g g > g g g g: 8160 -> 1371 jamp statements, 512 -> 274 kB, CPPProcess.cc compiles in 0.96s instead of 2.04s and the matrix elements come out 35% faster. |M|^2 agrees with the expanded output to 7e-16 (the optimisation reassociates the sums, exactly as the fortran one does), and the mg7 cross-section acceptance tests are unchanged. --jamp_optim=False recovers the expanded output. Co-Authored-By: Claude Opus 5 --- madgraph/iolibs/export_v4.py | 267 ++------------------- madgraph/iolibs/jamp_optimiser.py | 381 ++++++++++++++++++++++++++++++ madmatrix/model_handling.py | 229 ++++++++++++++++-- madmatrix/output.py | 4 +- 4 files changed, 614 insertions(+), 267 deletions(-) create mode 100644 madgraph/iolibs/jamp_optimiser.py diff --git a/madgraph/iolibs/export_v4.py b/madgraph/iolibs/export_v4.py index b99c7137f4..d774d7b0bd 100755 --- a/madgraph/iolibs/export_v4.py +++ b/madgraph/iolibs/export_v4.py @@ -49,6 +49,7 @@ import madgraph.iolibs.group_subprocs as group_subprocs import madgraph.iolibs.file_writers as writers import madgraph.iolibs.gen_infohtml as gen_infohtml +import madgraph.iolibs.jamp_optimiser as jamp_optimiser import madgraph.iolibs.template_files as template_files import madgraph.iolibs.ufo_expression_parsers as parsers import madgraph.iolibs.helas_call_writers as helas_call_writers @@ -223,7 +224,7 @@ def export_helas(self, HELAS_PATH): #=============================================================================== # ProcessExporterFortran #=============================================================================== -class ProcessExporterFortran(VirtualExporter): +class ProcessExporterFortran(VirtualExporter, jamp_optimiser.JampOptimiser): """Class to take care of exporting a set of matrix elements to Fortran (v4) format.""" @@ -233,15 +234,11 @@ class ProcessExporterFortran(VirtualExporter): 'output_options':{} } grouped_mode = False - jamp_optim = False - # how many times the JAMP optimisation called itself, for the record - myjamp_count = 0 # sum |M|^2 over one color flow per reversal pair instead of over every one # Folding the color matrix onto one line per reversal pair only works # where the template sums over NCOLORFOLD. get_color_data_lines is shared # by every fortran exporter, so this stays off unless the template agrees. jamp_fold = False - jamp_integer_walk = True # BLAS-3 for the color sum: all helicities at once as one right hand side. # None means take it when the library is there and the process is big # enough for it to pay. @@ -2820,12 +2817,10 @@ def get_JAMP_lines(self, col_amps, JAMP_format="JAMP(%s)", AMP_format="AMP(%s)", else: raise MadGraph5Error("Incorrect col_amps argument passed to get_JAMP_lines") - all_element = {} + # the coefficient matrix the optimisation below works on, built once + # from the same color amplitudes the expanded lines are written from + all_element = self.jamp_matrix(color_amplitudes) res_list = [] - # Every single amplitude carries a power of the number of colors in its - # coefficient, but a process only uses a handful of distinct powers, so - # build the corresponding fractions once instead of once per amplitude. - nc_powers = {} for i, coeff_list in enumerate(color_amplitudes): # It might happen that coeff_list is empty if this function was # called from get_JAMP_lines_split_order (i.e. if some color flow @@ -2859,16 +2854,6 @@ def get_JAMP_lines(self, col_amps, JAMP_format="JAMP(%s)", AMP_format="AMP(%s)", for (coefficient, amp_number) in coefs: if not coefficient: continue - try: - nc_power = nc_powers[coefficient[3]] - except KeyError: - nc_power = fractions.Fraction(3)**coefficient[3] - nc_powers[coefficient[3]] = nc_power - value = (1j if coefficient[2] else 1)* coefficient[0] * coefficient[1] * nc_power - if (i+1, amp_number) not in all_element: - all_element[(i+1, amp_number)] = value - else: - all_element[(i+1, amp_number)] += value if common_factor: res = (res + "%s" + AMP_format) % \ (self.coeff(coefficient[0], @@ -2887,51 +2872,25 @@ def get_JAMP_lines(self, col_amps, JAMP_format="JAMP(%s)", AMP_format="AMP(%s)", res = res + ')' res_list.append(res) - if 'jamp_optim' in self.cmd_options: - jamp_optim = banner_mod.ConfigFile.format_variable(self.cmd_options['jamp_optim'], bool, 'jamp_optim') - else: - # class default - jamp_optim = self.jamp_optim - - if not jamp_optim: + if not self.jamp_optim_enabled(): return res_list, 0 else: saved = list(res_list) - + if len(all_element) > 1000: logger.info("Computing Color-Flow optimization [%s term]", len(all_element)) start_time = time.time() - else: + else: start_time = 0 - + res_list = [] self.myjamp_count = 0 - # With one power of i shared by every coefficient, dividing it out - # leaves whole numbers to walk over -- they compare and hash exactly, - # and nothing has to be widened to complex. The phase goes back onto - # the JAMP coefficients afterwards, so the lines written are the same. - phase = self.jamp_global_phase(all_element) \ - if self.jamp_integer_walk else None - integral = False - if phase is not None: - whole = {} - for key, value in all_element.items(): - number = value / phase if phase != 1 else value - if isinstance(number, complex): - number = number.real - number = fractions.Fraction(number).limit_denominator(10**9) - if number.denominator != 1: - break - whole[key] = int(number) - else: - all_element.clear() - all_element.update(whole) - integral = True - if not integral: - phase = None - for key in all_element: - all_element[key] = complex(all_element[key]) + # The optimisation itself is language neutral (see jamp_optimiser); it + # is run one step at a time here rather than through + # optimise_jamp_matrix because the color basis symmetry has to be read + # off the matrix once the phase has been taken out of it. + phase = self.jamp_walk_integers(all_element) self.jamp_orbits = None # the color basis is read from the matrix element, which is not always # what is passed here: the split order version hands over one list of @@ -2940,14 +2899,10 @@ def get_JAMP_lines(self, col_amps, JAMP_format="JAMP(%s)", AMP_format="AMP(%s)", col_amps if symmetry_source is None else symmetry_source, all_element) if orbit and self.jamp_orbit else None new_mat, defs = self.optimise_jamp(all_element, symmetry=symmetry) - if phase is not None and phase != 1: - # the definitions hold ratios, which the phase cancels out of; only - # the coefficients on the JAMP lines carry it - for key in new_mat: - new_mat[key] = new_mat[key] * phase + self.jamp_apply_phase(new_mat, phase) if start_time: logger.info("Color-Flow passed to %s term in %ss. Introduce %i contraction", len(new_mat), int(time.time()-start_time), len(defs)) - + #misc.sprint("number of iteration", self.myjamp_count) def format(frac): @@ -3080,172 +3035,6 @@ def format(frac): return res_list, len(defs) - @staticmethod - def index_jamp_matrix(all_element, nb_col): - """Sorted lists of the positions of the non zero entries of the matrix, - by line and by column. An entry which is present but zero does not - count, and neither does a column outside the 0..nb_col range, so that - these indices list exactly the entries the plain scan would look at.""" - - lines = collections.defaultdict(list) - columns = collections.defaultdict(list) - for (i, j), value in all_element.items(): - if value and j < nb_col: - lines[i].append(j) - columns[j].append(i) - for line in lines.values(): - line.sort() - for column in columns.values(): - column.sort() - return lines, columns - - @staticmethod - def common_jamp_lines(columns, nb_line, j1, j2): - """Lines, in increasing order, where both columns j1 and j2 are non - zero. Both column lists are sorted, so this is a plain merge.""" - - left, right = columns.get(j1, []), columns.get(j2, []) - res = [] - pos1 = pos2 = 0 - while pos1 < len(left) and pos2 < len(right): - if left[pos1] == right[pos2]: - if left[pos1] < nb_line: - res.append(left[pos1]) - pos1 += 1 - pos2 += 1 - elif left[pos1] < right[pos2]: - pos1 += 1 - else: - pos2 += 1 - return res - - def optimise_jamp(self, all_element, nb_line=0, nb_col=0, added=0, - symmetry=None): - """ optimise problem of type Y = A X - A is a matrix (all_element) - X is the fortran name of the input. - The code iteratively add sub-expression jtemp[sub_add] - and recall itself (this is add to the X size) - - With a symmetry (see get_jamp_symmetry) the sub-expressions are - introduced by whole orbits of that symmetry instead of one at a - time, so that the result can be written as one recipe per orbit. - The orbits are then left in self.jamp_orbits. - """ - if symmetry: - return self.optimise_jamp_best(all_element, symmetry) - - self.myjamp_count +=1 - - if not nb_line: - for i,j in all_element: - if i+1 > nb_line: - nb_line = i+1 - if j+1> nb_col: - nb_col = j+1 - if nb_col > 600 and added==0: - all_element1, all_element2 = {}, {} - for (k1,k2) in all_element: - if k2 >= nb_col//2: - all_element2[(k1,1+k2-(nb_col//2))] = all_element[(k1,k2)] - else: - all_element1[(k1,k2)] = all_element[(k1,k2)] - - all_element1, newdef1 = self.optimise_jamp(all_element1) - nb_added1 = len(newdef1) - - all_element2, newdef2 = self.optimise_jamp(all_element2) - - for (k1,k2) in all_element2: - if k2 >= 0: - all_element1[(k1,k2+(nb_col//2)-1)] = all_element2[(k1,k2)] - if k2 < 0: - all_element1[(k1,k2-nb_added1)] = all_element2[(k1,k2)] - # new_def format: added,j1,j2,R, max_count - for k, j1,j2, R, c in newdef2: - if j2 > 0: - k2 = j2+nb_col//2 -1 - else: - k2 = j2-nb_added1 - if j1 > 0: - k1 = j1+nb_col//2 -1 - else: - k1 = j1-nb_added1 - newdef1.append((k+nb_added1, k1, k2, R, c)) - if newdef1: - all_element, new_def = self.optimise_jamp(all_element1, nb_line=0, nb_col=0, added=len(newdef1)) - newdef1 = newdef1 + new_def - return all_element, newdef1 - - # Index of the non zero entries, by line and by column. The matrix is - # very sparse (a color flow only gets a small share of the amplitudes) - # so walking the whole 0..nb_col range for every entry, as looking the - # columns up one by one in the matrix amounts to, spends nearly all of - # its time discovering zeros. - lines, columns = self.index_jamp_matrix(all_element, nb_col) - - max_count = 0 - all_index = [] - # how many lines have the same ratio between two given columns, keyed - # by the two columns and the ratio at once rather than by nested - # dictionaries: this is the innermost loop of the whole optimisation - operation = collections.defaultdict(int) - for (i,j1), v1 in all_element.items(): - line = lines.get(i) - if not line: - continue - for j2 in line[bisect.bisect_right(line, j1):]: - key = (j1, j2, all_element[(i,j2)]/v1) - operation[key] += 1 - count = operation[key] - if count > max_count: - max_count = count - all_index = [key] - elif count == max_count: - all_index.append(key) - - if max_count <= 1: - return all_element, [] - - to_add = [] - for index in all_index: - j1,j2,R = index - first = True - # only the lines where both columns are filled can contribute; the - # substitutions done here can empty some of them, so the values - # still have to be read back from the matrix - for i in self.common_jamp_lines(columns, nb_line, j1, j2): - v1 = all_element.get((i,j1), 0) - v2 = all_element.get((i,j2), 0) - if not v1 or not v2: - continue - if v2/v1 == R: - if first: - first = False - added +=1 - to_add.append((added,j1,j2,R, max_count)) - - all_element[(i,-added)] = v1 - del all_element[(i,j1)] #= 0 - del all_element[(i,j2)] #= 0 - - logger.log(5,"Define %d new shortcut reused %d times", len(to_add), max_count) - new_element, new_def = self.optimise_jamp(all_element, nb_line=nb_line, nb_col=nb_col, added=added) - for one_def in to_add: - new_def.insert(0, one_def) - return new_element, new_def - - - @staticmethod - def jamp_operation_count(new_mat, defs): - """Additions the result asks for: one per definition, plus what is left - in each line of the matrix.""" - - terms = collections.Counter() - for jamp, _var in new_mat: - terms[jamp] += 1 - return len(defs) + sum(max(0, count - 1) for count in terms.values()) - def optimise_jamp_best(self, all_element, symmetry): """Taking whole orbits only pays once there is enough of them to share: on a small matrix it can end up asking for more additions than the plain @@ -3415,30 +3204,6 @@ def blas_wanted(self, nfold): return True return nfold >= self.blas_min_ncolor - @staticmethod - def jamp_global_phase(all_element): - """The power of i every coefficient carries, when they all carry the - same one. A pure gluon process picks up one factor of i per f^abc, the - same for every term, so the whole matrix is real or wholly imaginary; - a quark line mixes the two and there is nothing to take out.""" - - phase = None - for value in all_element.values(): - if not value: - continue - number = complex(value) - if number.imag == 0: - here = 1 - elif number.real == 0: - here = 1j - else: - return None - if phase is None: - phase = here - elif phase != here: - return None - return phase - def get_jamp_folding(self, matrix_element): """Whether to sum |M|^2 over one line per reversal pair, and the (reverse, sign, representatives, slot) that goes with it. diff --git a/madgraph/iolibs/jamp_optimiser.py b/madgraph/iolibs/jamp_optimiser.py new file mode 100644 index 0000000000..f11a7a0208 --- /dev/null +++ b/madgraph/iolibs/jamp_optimiser.py @@ -0,0 +1,381 @@ +################################################################################ +# +# Copyright (c) 2009 The MadGraph5_aMC@NLO Development team and Contributors +# +# This file is a part of the MadGraph5_aMC@NLO project, an application which +# automatically generates Feynman diagrams and matrix elements for arbitrary +# high-energy processes in the Standard Model and beyond. +# +# It is subject to the MadGraph5_aMC@NLO license which should accompany this +# distribution. +# +# For more information, visit madgraph.phys.ucl.ac.be and amcatnlo.web.cern.ch +# +################################################################################ +"""Language neutral part of the color flow (JAMP) optimisation. + +Every backend writes the same object: the matrix of coefficients giving each +color flow as a combination of the amplitudes, + + JAMP(i) = sum_j A(i,j) * AMP(j) + +Written out as it stands that is one line per non zero entry, which for a +multi-gluon process is tens of thousands of them. The search below replaces the +repeated pieces by shared sub-expressions, so that the matrix is left with far +fewer entries and a list of definitions to compute first. + +Nothing here knows about fortran or C++: it takes the coefficient matrix and +gives back the reduced matrix and the definitions. The exporters print that in +their own language (get_JAMP_lines for fortran, get_jamp_accumulation_lines for +the C++/cudacpp writer). +""" + +from __future__ import absolute_import + +import bisect +import collections +import fractions +import logging +import time + +import madgraph.various.banner as banner_mod + +logger = logging.getLogger('madgraph.export_v4') + + +class JampOptimiser(object): + """The common sub-expression search over the JAMP coefficient matrix. + + Mixed into the exporters, which supply the printing. A subclass that hands + a symmetry to optimise_jamp must also provide optimise_jamp_best (only the + fortran exporter does, see export_v4).""" + + # Off by default: the plain output of a backend is the expanded one, and + # each exporter switches this on for itself. 'jamp_optim' in cmd_options + # (i.e. --jamp_optim=True|False at output time) wins over the class value. + jamp_optim = False + # how many times the JAMP optimisation called itself, for the record + myjamp_count = 0 + # take the power of i shared by every coefficient out before searching, so + # that the search walks over whole numbers (see optimise_jamp_matrix) + jamp_integer_walk = True + + def jamp_optim_enabled(self): + """Whether to run the optimisation, --jamp_optim first.""" + + cmd_options = getattr(self, 'cmd_options', None) or {} + if 'jamp_optim' in cmd_options: + return banner_mod.ConfigFile.format_variable( + cmd_options['jamp_optim'], bool, 'jamp_optim') + return self.jamp_optim + + @staticmethod + def jamp_matrix(color_amplitudes): + """The coefficient matrix all_element[(color flow, amplitude)] = value + of the color amplitudes, color flows numbered from 1 and amplitudes as + they number themselves. This is the input of the optimisation, and the + same value the expanded lines are written with.""" + + all_element = {} + # Every single amplitude carries a power of the number of colors in its + # coefficient, but a process only uses a handful of distinct powers, so + # build the corresponding fractions once instead of once per amplitude. + nc_powers = {} + for i, coeff_list in enumerate(color_amplitudes): + for (coefficient, amp_number) in coeff_list: + if not coefficient: + continue + try: + nc_power = nc_powers[coefficient[3]] + except KeyError: + nc_power = fractions.Fraction(3)**coefficient[3] + nc_powers[coefficient[3]] = nc_power + value = (1j if coefficient[2] else 1) * \ + coefficient[0] * coefficient[1] * nc_power + key = (i + 1, amp_number) + if key not in all_element: + all_element[key] = value + else: + all_element[key] += value + return all_element + + def jamp_walk_integers(self, all_element): + """Take the power of i shared by every coefficient out of the matrix + and return it, leaving all_element with whole numbers -- they compare + and hash exactly, and nothing has to be widened to complex. Returns + None when there is no such phase, all_element then being complex + throughout. The phase goes back on with jamp_apply_phase, so the lines + written are the same either way.""" + + phase = self.jamp_global_phase(all_element) \ + if self.jamp_integer_walk else None + if phase is not None: + whole = {} + for key, value in all_element.items(): + number = value / phase if phase != 1 else value + if isinstance(number, complex): + number = number.real + number = fractions.Fraction(number).limit_denominator(10**9) + if number.denominator != 1: + break + whole[key] = int(number) + else: + all_element.clear() + all_element.update(whole) + return phase + for key in all_element: + all_element[key] = complex(all_element[key]) + return None + + @staticmethod + def jamp_apply_phase(new_mat, phase): + """Put back the phase jamp_walk_integers took out. The definitions hold + ratios, which it cancels out of; only the coefficients left in the + matrix carry it.""" + + if phase is None or phase == 1: + return + for key in new_mat: + new_mat[key] = new_mat[key] * phase + + def optimise_jamp_matrix(self, all_element, symmetry=None): + """Run the optimisation over the coefficient matrix and return + (new_mat, defs): + - defs is a list of (i, op1, op2, frac, nb): definition number i is + op1 + frac*op2, where a positive operand is an amplitude and a + negative one is the definition number -op; + - new_mat is what the matrix is left with, keyed the same way as the + input except that a negative amplitude index means the definition + of that number. + + all_element is consumed (the optimisation works in place). The fortran + exporter runs the three steps itself, since it has to look at the + walked matrix to work out the color basis symmetry in between.""" + + if len(all_element) > 1000: + logger.info("Computing Color-Flow optimization [%s term]", + len(all_element)) + start_time = time.time() + else: + start_time = 0 + + self.myjamp_count = 0 + phase = self.jamp_walk_integers(all_element) + new_mat, defs = self.optimise_jamp(all_element, symmetry=symmetry) + self.jamp_apply_phase(new_mat, phase) + if start_time: + logger.info("Color-Flow passed to %s term in %ss. Introduce %i contraction", + len(new_mat), int(time.time()-start_time), len(defs)) + return new_mat, defs + + @staticmethod + def jamp_global_phase(all_element): + """The power of i every coefficient carries, when they all carry the + same one. A pure gluon process picks up one factor of i per f^abc, the + same for every term, so the whole matrix is real or wholly imaginary; + a quark line mixes the two and there is nothing to take out.""" + + phase = None + for value in all_element.values(): + if not value: + continue + number = complex(value) + if number.imag == 0: + here = 1 + elif number.real == 0: + here = 1j + else: + return None + if phase is None: + phase = here + elif phase != here: + return None + return phase + + @staticmethod + def index_jamp_matrix(all_element, nb_col): + """Sorted lists of the positions of the non zero entries of the matrix, + by line and by column. An entry which is present but zero does not + count, and neither does a column outside the 0..nb_col range, so that + these indices list exactly the entries the plain scan would look at.""" + + lines = collections.defaultdict(list) + columns = collections.defaultdict(list) + for (i, j), value in all_element.items(): + if value and j < nb_col: + lines[i].append(j) + columns[j].append(i) + for line in lines.values(): + line.sort() + for column in columns.values(): + column.sort() + return lines, columns + + @staticmethod + def common_jamp_lines(columns, nb_line, j1, j2): + """Lines, in increasing order, where both columns j1 and j2 are non + zero. Both column lists are sorted, so this is a plain merge.""" + + left, right = columns.get(j1, []), columns.get(j2, []) + res = [] + pos1 = pos2 = 0 + while pos1 < len(left) and pos2 < len(right): + if left[pos1] == right[pos2]: + if left[pos1] < nb_line: + res.append(left[pos1]) + pos1 += 1 + pos2 += 1 + elif left[pos1] < right[pos2]: + pos1 += 1 + else: + pos2 += 1 + return res + + def optimise_jamp(self, all_element, nb_line=0, nb_col=0, added=0, + symmetry=None): + """ optimise problem of type Y = A X + A is a matrix (all_element) + X is the fortran name of the input. + The code iteratively add sub-expression jtemp[sub_add] + and recall itself (this is add to the X size) + + With a symmetry (see get_jamp_symmetry) the sub-expressions are + introduced by whole orbits of that symmetry instead of one at a + time, so that the result can be written as one recipe per orbit. + The orbits are then left in self.jamp_orbits. + """ + if symmetry: + return self.optimise_jamp_best(all_element, symmetry) + + self.myjamp_count +=1 + + if not nb_line: + for i,j in all_element: + if i+1 > nb_line: + nb_line = i+1 + if j+1> nb_col: + nb_col = j+1 + if nb_col > 600 and added==0: + all_element1, all_element2 = {}, {} + for (k1,k2) in all_element: + if k2 >= nb_col//2: + all_element2[(k1,1+k2-(nb_col//2))] = all_element[(k1,k2)] + else: + all_element1[(k1,k2)] = all_element[(k1,k2)] + + all_element1, newdef1 = self.optimise_jamp(all_element1) + nb_added1 = len(newdef1) + + all_element2, newdef2 = self.optimise_jamp(all_element2) + + for (k1,k2) in all_element2: + if k2 >= 0: + all_element1[(k1,k2+(nb_col//2)-1)] = all_element2[(k1,k2)] + if k2 < 0: + all_element1[(k1,k2-nb_added1)] = all_element2[(k1,k2)] + # new_def format: added,j1,j2,R, max_count + for k, j1,j2, R, c in newdef2: + if j2 > 0: + k2 = j2+nb_col//2 -1 + else: + k2 = j2-nb_added1 + if j1 > 0: + k1 = j1+nb_col//2 -1 + else: + k1 = j1-nb_added1 + newdef1.append((k+nb_added1, k1, k2, R, c)) + if newdef1: + all_element, new_def = self.optimise_jamp(all_element1, nb_line=0, nb_col=0, added=len(newdef1)) + newdef1 = newdef1 + new_def + return all_element, newdef1 + + # Index of the non zero entries, by line and by column. The matrix is + # very sparse (a color flow only gets a small share of the amplitudes) + # so walking the whole 0..nb_col range for every entry, as looking the + # columns up one by one in the matrix amounts to, spends nearly all of + # its time discovering zeros. + lines, columns = self.index_jamp_matrix(all_element, nb_col) + + max_count = 0 + all_index = [] + # how many lines have the same ratio between two given columns, keyed + # by the two columns and the ratio at once rather than by nested + # dictionaries: this is the innermost loop of the whole optimisation + operation = collections.defaultdict(int) + for (i,j1), v1 in all_element.items(): + line = lines.get(i) + if not line: + continue + for j2 in line[bisect.bisect_right(line, j1):]: + key = (j1, j2, all_element[(i,j2)]/v1) + operation[key] += 1 + count = operation[key] + if count > max_count: + max_count = count + all_index = [key] + elif count == max_count: + all_index.append(key) + + if max_count <= 1: + return all_element, [] + + to_add = [] + for index in all_index: + j1,j2,R = index + first = True + # only the lines where both columns are filled can contribute; the + # substitutions done here can empty some of them, so the values + # still have to be read back from the matrix + for i in self.common_jamp_lines(columns, nb_line, j1, j2): + v1 = all_element.get((i,j1), 0) + v2 = all_element.get((i,j2), 0) + if not v1 or not v2: + continue + if v2/v1 == R: + if first: + first = False + added +=1 + to_add.append((added,j1,j2,R, max_count)) + + all_element[(i,-added)] = v1 + del all_element[(i,j1)] #= 0 + del all_element[(i,j2)] #= 0 + + logger.log(5,"Define %d new shortcut reused %d times", len(to_add), max_count) + new_element, new_def = self.optimise_jamp(all_element, nb_line=nb_line, nb_col=nb_col, added=added) + for one_def in to_add: + new_def.insert(0, one_def) + return new_element, new_def + + @staticmethod + def jamp_operation_count(new_mat, defs): + """Additions the result asks for: one per definition, plus what is left + in each line of the matrix.""" + + terms = collections.Counter() + for jamp, _var in new_mat: + terms[jamp] += 1 + return len(defs) + sum(max(0, count - 1) for count in terms.values()) + + @staticmethod + def jamp_definition_order(defs): + """The definitions in the order they can be computed while the + amplitudes are produced one at a time, together with the amplitude each + one is ready after. + + Returns (order, ready): order lists the definition numbers, ready maps + a definition number onto the last amplitude it needs (transitively). + Sorting by that amplitude keeps the list topological, since a + definition never needs fewer amplitudes than the ones it is built + from.""" + + ready = {} + rank = {} + for position, (i, amp1, amp2, _frac, _nb) in enumerate(defs): + last = 0 + for amp in (amp1, amp2): + last = max(last, amp if amp > 0 else ready[-amp]) + ready[i] = last + rank[i] = position + order = sorted(ready, key=lambda i: (ready[i], rank[i])) + return order, ready diff --git a/madmatrix/model_handling.py b/madmatrix/model_handling.py index b4d1bac787..4c508f71a4 100644 --- a/madmatrix/model_handling.py +++ b/madmatrix/model_handling.py @@ -18,6 +18,7 @@ from madgraph.iolibs import export_cpp, export_mg7 from madgraph.iolibs import file_writers as writers +from madgraph.iolibs import jamp_optimiser import aloha from aloha import aloha_writers @@ -1904,8 +1905,16 @@ 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[ncolor] = {}; // all zeros (NB: vector cxtype_v IS initialized to 0, but scalar cxtype is NOT, if "= {}" is missing!)""") + # 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. + nb_tmp_jamp = getattr(self.helas_call_writer, 'nb_tmp_jamp', 0) + if nb_tmp_jamp: + ret_lines.append(""" + // [jampTmp: partial sums of amplitudes that several color flows share, so that they are computed only once] + cxtype_sv jampTmp_sv[%i];""" % nb_tmp_jamp) + ret_lines.append(""" // === Calculate wavefunctions and amplitudes for all diagrams in all processes === // === (for one event in CUDA, for one - or two in mixed mode - SIMD event pages in C++ === @@ -2234,13 +2243,18 @@ def get_reset_jamp_lines(self, color_amplitudes): # AV - define a custom HelasCallWriter # (NB: enable this via ProcessExporterMadMatrix.helas_exporter in output.py - this fixes #341) -class MadMatrixUFOHelasCallWriter(helas_call_writers.GPUFOHelasCallWriter): +class MadMatrixUFOHelasCallWriter(helas_call_writers.GPUFOHelasCallWriter, + jamp_optimiser.JampOptimiser): """ A Custom HelasCallWriter """ # Flavor-mask optimization: skip wavefunction/amplitude calls that vanish # for the selected flavor (see super_get_matrix_element_calls). Toggled by # the output command's --mask=True|False; default on. use_flavor_mask = True + # Write the color flows through the shared sub-expressions the color-flow + # optimisation finds, instead of one line per (color flow, amplitude) pair + # (see build_jamp_plan). Toggled by --jamp_optim=True|False. + jamp_optim = True # Class structure information # - object # - dict(object) [built-in] @@ -2249,8 +2263,9 @@ class MadMatrixUFOHelasCallWriter(helas_call_writers.GPUFOHelasCallWriter): # - UFOHelasCallWriter(HelasCallWriter) [in madgraph/iolibs/helas_call_writers.py] # - CPPUFOHelasCallWriter(UFOHelasCallWriter) [in madgraph/iolibs/helas_call_writers.py] # - GPUFOHelasCallWriter(CPPUFOHelasCallWriter) [in madgraph/iolibs/helas_call_writers.py] - # - MadMatrixUFOHelasCallWriter(GPUFOHelasCallWriter) - # This class + # - MadMatrixUFOHelasCallWriter(GPUFOHelasCallWriter, JampOptimiser) + # This class (JampOptimiser is in madgraph/iolibs/jamp_optimiser.py and + # brings the color-flow optimisation shared with the fortran exporter) def __init__(self, *args, **opts): @@ -2404,6 +2419,165 @@ def format_coupling(self, call): def format_call(call): return call.replace('(','( ').replace(')',' )').replace(',',', ') + # --- Color flows through shared sub-expressions -------------------------- + # + # Written out as it stands, a color flow is one 'jamp_sv[i] += c*amp_sv[0]' + # per (color flow, amplitude) pair: eight thousand of them for g g > g g g + # g. The optimisation in jamp_optimiser finds the partial sums that several + # flows have in common and returns them as definitions + # + # TMP(i) = + frac * + # + # where an operand is either an amplitude or an earlier definition, leaving + # the color flows as a much shorter combination of those definitions. + # + # The fortran output has AMP(NGRAPHS) to read the amplitudes back from, so + # it prints the definitions as they come. Here every amplitude passes + # through the single slot amp_sv[0] and is gone by the next diagram, so the + # definitions are accumulated instead: each amplitude is added into the one + # definition that uses it while it is still there, and the definitions of + # definitions follow as soon as everything they need has been seen (see + # jamp_definition_order). + + @staticmethod + def jamp_number(value): + """A C++ literal for a real coefficient, kept exact where it can be + (an integer, or a ratio the compiler divides out itself).""" + + frac = Fraction(value).limit_denominator(10**9) + if float(frac) == float(value): + if frac.denominator == 1: + return '%d.' % frac.numerator + return '%d. / %d.' % (frac.numerator, frac.denominator) + text = '%.17g' % value + if '.' not in text and 'e' not in text and 'n' not in text: + text += '.' + return text + + @classmethod + def jamp_factor(cls, value): + """(sign, factor) of a JAMP coefficient, the factor being the C++ text + multiplying the operand (empty when the coefficient is +-1).""" + + number = complex(value) + if number.imag == 0: + magnitude, imaginary = number.real, False + elif number.real == 0: + magnitude, imaginary = number.imag, True + else: + # never seen in practice: a color coefficient is real or imaginary + return 1, 'cxtype( %s, %s ) * ' % (cls.jamp_number(number.real), + cls.jamp_number(number.imag)) + sign = -1 if magnitude < 0 else 1 + magnitude = abs(magnitude) + if magnitude == 1: + return sign, ('cxtype( 0, 1 ) * ' if imaginary else '') + if imaginary: + return sign, '%s * cxtype( 0, 1 ) * ' % cls.jamp_number(magnitude) + return sign, '%s * ' % cls.jamp_number(magnitude) + + @classmethod + def jamp_statement(cls, target, terms, assign): + """'target = t1 - t2;' (assign) or 'target += t1 - t2;', from a list of + (coefficient, operand) terms.""" + + pieces = [] + for pos, (value, name) in enumerate(terms): + sign, factor = cls.jamp_factor(value) + if pos == 0 and not assign and len(terms) == 1: + # the common case: keep the sign on the operator, as the + # expanded output does + return '%s %s= %s%s;' % (target, '-' if sign < 0 else '+', + factor, name) + if pos == 0: + pieces.append('%s%s%s' % ('-' if sign < 0 else '', factor, name)) + else: + pieces.append('%s %s%s' % ('-' if sign < 0 else '+', factor, name)) + return '%s %s %s;' % (target, '=' if assign else '+=', ' '.join(pieces)) + + def build_jamp_plan(self, color_amplitudes): + """Work out how the color flows are built from shared sub-expressions, + and return (ntmp, captures, combines, final): + - captures[n] are the lines to write while amplitude n sits in + amp_sv[0], as (line, target, is_first_write) so that a masked + amplitude can be told to zero its target first; + - combines[n] are the definitions ready once amplitude n has been + added, to write just after it; + - final are the lines assembling jamp_sv out of the definitions. + Returns None when there is nothing to share, so that the caller keeps + the expanded output.""" + + if not self.jamp_optim_enabled(): + return None + all_element = self.jamp_matrix(color_amplitudes) + if not all_element: + return None + new_mat, defs = self.optimise_jamp_matrix(all_element) + if not defs: + return None + order, ready = self.jamp_definition_order(defs) + definition = {i: (amp1, amp2, frac) for i, amp1, amp2, frac, _nb in defs} + + # what each amplitude has to be added into while it is still in amp_sv + captures = defaultdict(list) # amplitude -> [(target, coefficient)] + for i, amp1, amp2, frac, _nb in defs: + for amp, coefficient in ((amp1, 1), (amp2, frac)): + if amp > 0: + captures[amp].append(('jampTmp_sv[%d]' % (i - 1), coefficient)) + for (jamp, var), factor in sorted(new_mat.items()): + if var > 0 and factor: + captures[var].append(('jamp_sv[%d]' % (jamp - 1), factor)) + + # the definitions in the order they become available, grouped by the + # amplitude they are ready after + ready_after = defaultdict(list) + for i in order: + ready_after[ready[i]].append(i) + + started = set() # definitions already assigned to + done = set() # definitions holding their full value + capture_lines = defaultdict(list) + combine_lines = defaultdict(list) + for amp in sorted(set(list(captures) + list(ready_after))): + for target, coefficient in captures.get(amp, []): + # jamp_sv is zeroed at the top of the event page, so it is only + # the definitions that have to start with an assignment + first = target.startswith('jampTmp_sv') and target not in started + if first: + started.add(target) + capture_lines[amp].append( + (self.jamp_statement(target, [(coefficient, 'amp_sv[0]')], + first), target, first)) + for i in ready_after.get(amp, []): + amp1, amp2, frac = definition[i] + operands = [operand for operand in (amp1, amp2) if operand < 0] + assert all(-operand in done for operand in operands), \ + 'a color-flow definition is used before it is complete' + done.add(i) + if not operands: + continue # both operands were amplitudes, already added + terms = [(coefficient, 'jampTmp_sv[%d]' % (-operand - 1)) + for operand, coefficient in ((amp1, 1), (amp2, frac)) + if operand < 0] + target = 'jampTmp_sv[%d]' % (i - 1) + first = target not in started + started.add(target) + combine_lines[amp].append( + self.jamp_statement(target, terms, first)) + assert len(started) == len(defs) == len(done), \ + 'a color-flow definition is never written' + + # what is left of the color flows: a combination of the definitions + final = [] + by_jamp = defaultdict(list) + for (jamp, var), factor in sorted(new_mat.items()): + if var < 0 and factor: + by_jamp[jamp].append((factor, 'jampTmp_sv[%d]' % (-var - 1))) + for jamp in sorted(by_jamp): + final.append(self.jamp_statement('jamp_sv[%d]' % (jamp - 1), + by_jamp[jamp], False)) + return len(defs), capture_lines, combine_lines, final + # AV - replace helas_call_writers.GPUFOHelasCallWriter method (improve formatting) def super_get_matrix_element_calls(self, matrix_element, color_amplitudes, multi_channel_map): """Return a list of strings, corresponding to the Helas calls for the matrix element""" @@ -2412,6 +2586,7 @@ def super_get_matrix_element_calls(self, matrix_element, color_amplitudes, multi assert isinstance(matrix_element, helas_objects.HelasMatrixElement), \ '%s not valid argument for get_matrix_element_calls' % \ type(matrix_element) + self.nb_tmp_jamp = 0 # Do not reuse the wavefunctions for loop matrix elements if isinstance(matrix_element, loop_helas_objects.LoopHelasMatrixElement): return self.get_loop_matrix_element_calls(matrix_element) @@ -2422,6 +2597,12 @@ def super_get_matrix_element_calls(self, matrix_element, color_amplitudes, multi if namp not in color: color[namp] = {} color[namp][njamp] = coeff + # Color flows through shared sub-expressions (None to write them out + # one (color flow, amplitude) pair at a time, as before) + jamp_plan = self.build_jamp_plan(color_amplitudes) + self.nb_tmp_jamp = jamp_plan[0] if jamp_plan else 0 + if jamp_plan is not None: + _ntmp, jamp_captures, jamp_combines, jamp_final = jamp_plan me = matrix_element.get('diagrams') matrix_element.reuse_outdated_wavefunctions(me) ###misc.sprint(multi_channel_map) @@ -2563,27 +2744,45 @@ def _guard_open(group_mask): amp_block.append(" numerators_sv[%i] += cxabs2( amp_sv[0] );" % (diagnum-1)) amp_block.append(" denominators_sv += cxabs2( amp_sv[0] );") amp_block.append("}") - for njamp, coeff in color[namp].items(): - scoeff = OneProcessExporterMadMatrix.coeff(*coeff) # AV - if scoeff[0] == '+' : scoeff = scoeff[1:] - scoeff = scoeff.replace('(','( ') - scoeff = scoeff.replace(')',' )') - scoeff = scoeff.replace(',',', ') - scoeff = scoeff.replace('*',' * ') - scoeff = scoeff.replace('/',' / ') - 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 # The amplitude (and the jamp/channel contributions that read its # amp_sv[0]) only contributes for the flavors in the diagram's # mask, so guard the whole block as a unit. gmask = diag_group_mask.get(id(diagram)) + before_guard = [] + if jamp_plan is None: + for njamp, coeff in color[namp].items(): + scoeff = OneProcessExporterMadMatrix.coeff(*coeff) # AV + if scoeff[0] == '+' : scoeff = scoeff[1:] + scoeff = scoeff.replace('(','( ') + scoeff = scoeff.replace(')',' )') + scoeff = scoeff.replace(',',', ') + scoeff = scoeff.replace('*',' * ') + scoeff = scoeff.replace('/',' / ') + 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: + for line, target, first in jamp_captures.get(namp, []): + if first and gmask is not None: + # the guard can skip the line that opens this + # sub-expression, so it has to start from zero + before_guard.append('%s = cxzero_sv();' % target) + line = line.replace(' = ', ' += ', 1) + amp_block.append(line) + res.extend(before_guard) if gmask is not None: res.append(_guard_open(gmask)) res.extend(amp_block) res.append('}') else: res.extend(amp_block) + if jamp_plan is not None: + # sub-expressions that have now seen every amplitude they + # need: they always run, whatever the flavor + res.extend(jamp_combines.get(namp, [])) if len(diagram.get('amplitudes')) == 0 : res.append('// (none)') # AV + if jamp_plan is not None: + res.append('\n // *** COLOR FLOWS FROM THE SHARED SUB-EXPRESSIONS ***') + res.extend(jamp_final) ###res.append('\n // *** END OF DIAGRAMS ***' ) # AV - no longer needed ('COLOR MATRIX BELOW') return res diff --git a/madmatrix/output.py b/madmatrix/output.py index 1820b10e96..339b488177 100644 --- a/madmatrix/output.py +++ b/madmatrix/output.py @@ -177,9 +177,11 @@ def generate_subprocess_directory(self, matrix_element, cpp_helas_call_writer, p misc.sprint(' type(proc_number)=%s me=%s'%(type(proc_number) if proc_number is not None else None, proc_number)) # e.g. int misc.sprint("need to link", self.to_link_in_P) # Propagate the --mask toggle to the helas call writer that emits the - # guarded wavefunction/amplitude calls. + # guarded wavefunction/amplitude calls, and the output command line as + # a whole for the --jamp_optim toggle of the color-flow optimisation. if cpp_helas_call_writer is not None: cpp_helas_call_writer.use_flavor_mask = self.use_flavor_mask + cpp_helas_call_writer.cmd_options = self.opt.get('output_options', {}) out = super().generate_subprocess_directory(matrix_element, cpp_helas_call_writer, proc_number) return out From 332bf68df88427a82c20c2437763b359f65ce55d Mon Sep 17 00:00:00 2001 From: Olivier Mattelaer Date: Fri, 7 Aug 2026 23:20:11 +0200 Subject: [PATCH 37/42] fix the FPTYPE=f build of the madmatrix check_sa driver run_matrix_mode passed the fptype copy of the masses to classic_rambo::get_momenta, whose signature takes std::vector. At FPTYPE=d/m fptype is double so this compiled by accident; at FPTYPE=f it is float and there is no vector-to-vector conversion, so any generated process failed to build in single precision. Pass massesD, the double vector the masses are read into, instead. Both vectors are already in scope and both are still needed: the fptype copy keeps feeding RamboSamplingKernelHost. Widening the get_momenta signature would have been wrong here, since classic_rambo works in double throughout precisely so it reproduces the phase-space point of the Fortran/C++ 'check' drivers -- feeding it float-rounded masses would perturb that reference. Checked on g g > g g g with a clean rebuild at each precision: all three build, the phase-space point is identical, and the matrix elements agree to ~2 ulp of float (f vs d 2.7e-7 relative, m vs d 1.2e-8). Co-Authored-By: Claude Opus 5 --- madgraph/iolibs/template_files/madmatrix/check_sa.cc | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/madgraph/iolibs/template_files/madmatrix/check_sa.cc b/madgraph/iolibs/template_files/madmatrix/check_sa.cc index 68e93edb50..e7a36ca65d 100644 --- a/madgraph/iolibs/template_files/madmatrix/check_sa.cc +++ b/madgraph/iolibs/template_files/madmatrix/check_sa.cc @@ -648,8 +648,10 @@ namespace } const std::vector masses( massesD.begin(), massesD.end() ); + // NB: feed the double-precision masses to the classic RAMBO, which works in + // double throughout: 'masses' is fptype and would not convert at FPTYPE=f. std::vector> point = - classic_rambo::get_momenta( CPPProcess::npari, (double)kEnergy, masses, rambowgt ); + classic_rambo::get_momenta( CPPProcess::npari, (double)kEnergy, massesD, rambowgt ); // alpha_s from the param card so the couplings match the Fortran/C++ // 'check' drivers (UMAMI otherwise falls back to a hardcoded g_s). From 20446e4b290fd29ff2c7c6ba66f550cff828b061 Mon Sep 17 00:00:00 2001 From: Olivier Mattelaer Date: Fri, 7 Aug 2026 23:22:41 +0200 Subject: [PATCH 38/42] send the C++ color sum through a host BLAS The cuBLAS path in color_sum.cc was GPU only: the C++ side had a hand written triangular SIMD loop and nothing else. The color matrix does not depend on the helicity, so on C++ too the jamps of every good helicity of one event page are the columns of one right hand side and the whole color sum becomes two SYMM calls, real and imaginary parts apart. Measured first. On the trace basis the C++ color sum is a real share of the matrix element: 31% for g g > g g g g and 14% for g g > t t~ g g g at ncolor=120 (sampling and a gutted-color-sum difference agree), against 5% to 11% at ncolor=24. Accelerate's SYMM does about twice the flops of the triangular loop and still wins by 9x on its own, so the color sum drops to 7% of the matrix element and the whole thing runs 1.30x faster on g g > g g g g, 1.13x on g g > t t~ g g g. Taken exactly where the Fortran color sum takes it - the same DSYMM probe, the same ncolor threshold - so with BLAS off, or below the threshold, color_sum.cc and CPPProcess.cc are character for character the files written before any of this existed. Above it both paths are written out and CPPBLAS=hasNoBlas still builds the old one. The batched call bypasses nothing: the jamp2 sums for the color selection, the numerators and the denominators all stay inside calculate_jamps and still happen once per helicity. What color_sum_cpu did on top of the sum was feed MEs_ighel, the running sum over helicities the event by event choice of helicity reads, so the batched call rebuilds those itself. Mixed precision is supported rather than switched off, through SSYMM. |M|^2 agrees with a non-BLAS build to one ulp (some points bit identical); the summation order differs, so it is not bit for bit. The C++ to Fortran gap at the same phase space point is 6.5e-15, forty times larger. Co-Authored-By: Claude Opus 5 --- madgraph/iolibs/export_cpp.py | 12 +- madgraph/iolibs/export_v4.py | 11 ++ .../template_files/madmatrix/color_sum.cc | 2 +- .../template_files/madmatrix/color_sum.h | 21 +++ .../madmatrix/color_sum_blas.inc | 166 ++++++++++++++++++ .../madmatrix/color_sum_blas_loop.inc | 23 +++ .../template_files/madmatrix/madmatrix.mk | 35 +++- .../madmatrix/madmatrix_standalone.mk | 2 +- .../madmatrix/process_sigmaKin_function.inc | 4 +- madmatrix/model_handling.py | 33 ++++ madmatrix/output.py | 19 +- 11 files changed, 315 insertions(+), 13 deletions(-) create mode 100644 madgraph/iolibs/template_files/madmatrix/color_sum_blas.inc create mode 100644 madgraph/iolibs/template_files/madmatrix/color_sum_blas_loop.inc diff --git a/madgraph/iolibs/export_cpp.py b/madgraph/iolibs/export_cpp.py index 354815b369..87c9bc0fc7 100755 --- a/madgraph/iolibs/export_cpp.py +++ b/madgraph/iolibs/export_cpp.py @@ -2688,17 +2688,21 @@ def copy_template(self, model): if self.template_src_make: # Copy src Makefile makefile = self.read_template_file(self.template_src_make) % \ - {'model': self.get_model_name(model.get('name')), - 'cpp_compiler': self.opt['cpp_compiler'] if self.opt['cpp_compiler'] else 'g++'} + self.get_makefile_replace_dict(model) open(os.path.join('src', 'Makefile'), 'w').write(makefile) if self.template_Sub_make: # Copy SubProcesses Makefile makefile = self.read_template_file(self.template_Sub_make) % \ - {'model': self.get_model_name(model.get('name')), - 'cpp_compiler': self.opt['cpp_compiler'] if self.opt['cpp_compiler'] else 'g++'} + self.get_makefile_replace_dict(model) open(os.path.join('SubProcesses', 'Makefile'), 'w').write(makefile) + def get_makefile_replace_dict(self, model): + """Template replacements for the src and SubProcesses makefiles.""" + + return {'model': self.get_model_name(model.get('name')), + 'cpp_compiler': self.opt['cpp_compiler'] if self.opt['cpp_compiler'] else 'g++'} + #=========================================================================== # Helper functions #=========================================================================== diff --git a/madgraph/iolibs/export_v4.py b/madgraph/iolibs/export_v4.py index b99c7137f4..e3842aa9b9 100755 --- a/madgraph/iolibs/export_v4.py +++ b/madgraph/iolibs/export_v4.py @@ -3396,6 +3396,17 @@ def blas_is_available(cls): shutil.rmtree(work, ignore_errors=True) return cls._blas_available + @classmethod + def blas_available_flags(cls): + """What a BLAS carrying DSYMM needs on the link line, empty when there + is none. Unlike blas_link_flags this does not ask whether BLAS was + wanted, only whether it is there, which is what a backend deciding for + itself (the C++ color sum) needs.""" + + if not cls.blas_is_available(): + return '' + return cls._blas_flags + def blas_link_flags(self): """What to link the color sum against, empty when BLAS is not taken.""" diff --git a/madgraph/iolibs/template_files/madmatrix/color_sum.cc b/madgraph/iolibs/template_files/madmatrix/color_sum.cc index 30c6799932..d0805497f5 100644 --- a/madgraph/iolibs/template_files/madmatrix/color_sum.cc +++ b/madgraph/iolibs/template_files/madmatrix/color_sum.cc @@ -151,7 +151,7 @@ namespace mg5amcCpu #endif } #endif - +%(cpp_blas_color_sum)s //-------------------------------------------------------------------------- #ifdef MGONGPUCPP_GPUIMPL diff --git a/madgraph/iolibs/template_files/madmatrix/color_sum.h b/madgraph/iolibs/template_files/madmatrix/color_sum.h index 347184c4e1..b405cbad27 100644 --- a/madgraph/iolibs/template_files/madmatrix/color_sum.h +++ b/madgraph/iolibs/template_files/madmatrix/color_sum.h @@ -14,6 +14,10 @@ #include "CPPProcess.h" #include "GpuAbstraction.h" +#ifdef MGONGPU_CPP_HAS_BLAS +#include // the batched C++ color sum keeps the jamps of every good helicity +#endif + #ifdef MGONGPUCPP_GPUIMPL namespace mg5amcGpu #else @@ -76,6 +80,23 @@ namespace mg5amcCpu //-------------------------------------------------------------------------- + // Only defined for processes whose color matrix is large enough that the + // BLAS call is worth setting up (see blas_wanted): the color sum for every + // good helicity of one event page in one go. +#ifndef MGONGPUCPP_GPUIMPL +#ifdef MGONGPU_CPP_HAS_BLAS + void + color_sum_cpu_blas( fptype* allMEs, // input/output: allMEs[nevt], add |M|^2 summed over all good helicities + fptype_sv* MEs_ighel, // output: [ncomb] running sum of |M|^2 up to ighel (first - and/or only - neppV page) + fptype_sv* MEs_ighel2, // output: [ncomb] the same for the second neppV page (mixed mode only) + const cxtype_sv* ghelAllJamp_sv, // input: jamp_sv[nGoodHel][nParity*ncolor] for all good helicities + const int nGoodHel, // input: number of good helicities + const int ievt0 ); // input: first event number in current C++ event page +#endif +#endif + + //-------------------------------------------------------------------------- + #ifdef MGONGPUCPP_GPUIMPL void color_sum_gpu( fptype* ghelAllMEs, // output: allMEs super-buffer for nGoodHel <= ncomb individual helicities (index is ighel) diff --git a/madgraph/iolibs/template_files/madmatrix/color_sum_blas.inc b/madgraph/iolibs/template_files/madmatrix/color_sum_blas.inc new file mode 100644 index 0000000000..f32e89acbb --- /dev/null +++ b/madgraph/iolibs/template_files/madmatrix/color_sum_blas.inc @@ -0,0 +1,166 @@ +// Copyright (C) 2020-2026 CERN and UCLouvain. +// Licensed under the GNU Lesser General Public License (version 3 or later). +// Integrated with the MadGraph7 project in Feb 2026. +// +// The C++ color sum for every good helicity at once, through BLAS. This is +// only written out for processes whose color matrix is large enough that the +// call is worth setting up (see blas_wanted); everywhere else color_sum.cc is +// character for character the file written before any of this existed. + + //-------------------------------------------------------------------------- + +#ifndef MGONGPUCPP_GPUIMPL +#ifdef MGONGPU_CPP_HAS_BLAS + +#if defined MGONGPU_CPPSIMD and defined MGONGPU_FPTYPE_DOUBLE and defined MGONGPU_FPTYPE2_FLOAT + constexpr int nParityCS = 2; // mixed mode merges two neppV pages into one call +#else + constexpr int nParityCS = 1; +#endif + + // The color matrix does not depend on the helicity, so the jamps of every + // good helicity (and of every event in the SIMD page) are columns of one + // right hand side and the whole color sum becomes two SYMM calls. SYMM is + // real, so the real and the imaginary part of JAMP go through separately, + // which is the same property of M being real that the scalar sum uses to + // rewrite (A-iB)M(A+iB) as AMA + BMB (see #475). + // + // The scalar sum walks the upper triangle with its off diagonal doubled; + // SYMM wants the whole symmetric matrix with each entry counted once. What + // is written out is normalized row by row by colorDenom[icol], which is not + // symmetric when the denominators differ; only the symmetric part of a + // matrix contributes to a quadratic form, so take it explicitly. When all + // the denominators agree (the usual case) this is exactly colorMatrix/denom. + struct SymmetricNormalizedColorMatrix + { + constexpr SymmetricNormalizedColorMatrix() + : value() + { + for( int icol = 0; icol < ncolor; icol++ ) + for( int jcol = 0; jcol < ncolor; jcol++ ) + value[icol * ncolor + jcol] = + ( colorMatrix[icol][jcol] / colorDenom[icol] + colorMatrix[jcol][icol] / colorDenom[jcol] ) / 2; + } + fptype2 value[ncolor * ncolor]; + }; + + // The Fortran BLAS interface, which every implementation exports (the + // reference BLAS ships no CBLAS of its own). Column major, as in the + // Fortran color sum. + extern "C" + { + void dsymm_( const char* side, const char* uplo, const int* m, const int* n, + const double* alpha, const double* a, const int* lda, + const double* b, const int* ldb, + const double* beta, double* c, const int* ldc ); + void ssymm_( const char* side, const char* uplo, const int* m, const int* n, + const float* alpha, const float* a, const int* lda, + const float* b, const int* ldb, + const float* beta, float* c, const int* ldc ); + } + + static inline void + blas_symm( const int m, const int n, const double* a, const double* b, double* c ) + { + const char side = 'L', uplo = 'U'; + const double alpha = 1, beta = 0; + dsymm_( &side, &uplo, &m, &n, &alpha, a, &m, b, &m, &beta, c, &m ); + } + + static inline void + blas_symm( const int m, const int n, const float* a, const float* b, float* c ) + { + const char side = 'L', uplo = 'U'; + const float alpha = 1, beta = 0; + ssymm_( &side, &uplo, &m, &n, &alpha, a, &m, b, &m, &beta, c, &m ); + } + + void + color_sum_cpu_blas( fptype* allMEs, // input/output: allMEs[nevt], add |M|^2 summed over all good helicities + fptype_sv* MEs_ighel, // output: [ncomb] running sum of |M|^2 up to ighel (first - and/or only - neppV page) + fptype_sv* MEs_ighel2, // output: [ncomb] the same for the second neppV page (mixed mode only) + const cxtype_sv* ghelAllJamp_sv, // input: jamp_sv[nGoodHel][nParity*ncolor] for all good helicities + const int nGoodHel, // input: number of good helicities + const int ievt0 ) // input: first event number in current C++ event page + { + static constexpr auto cfsym = SymmetricNormalizedColorMatrix(); + constexpr int nevtB = nParityCS * neppV; // events covered by one call + const int ncol = nGoodHel * nevtB; // number of BLAS right hand side columns + // Column major scratch: JR/JI hold the ncolor x ncol jamps, ZR/ZI take the + // SYMM results and MEcol one |M|^2 per column. Kept on the heap and grown + // once per thread: for ncolor=120 and ncomb=128 this is a few hundred kB. + static thread_local std::vector scratch; + const size_t need = 4 * (size_t)ncolor * ncol + ncol; + if( scratch.size() < need ) scratch.resize( need ); + fptype2* JR = scratch.data(); + fptype2* JI = JR + (size_t)ncolor * ncol; + fptype2* ZR = JI + (size_t)ncolor * ncol; + fptype2* ZI = ZR + (size_t)ncolor * ncol; + fptype2* MEcol = ZI + (size_t)ncolor * ncol; + // Transpose the jamps into the column major right hand side: colour is the + // fast index, (helicity, event) the slow one. + for( int ighel = 0; ighel < nGoodHel; ighel++ ) + { + const cxtype_sv* jamp_sv = ghelAllJamp_sv + (size_t)ighel * nParityCS * ncolor; + for( int ip = 0; ip < nParityCS; ip++ ) + for( int ieppV = 0; ieppV < neppV; ieppV++ ) + { + const size_t off = (size_t)( ighel * nevtB + ip * neppV + ieppV ) * ncolor; + for( int icol = 0; icol < ncolor; icol++ ) + { +#ifdef MGONGPU_CPPSIMD + JR[off + icol] = cxreal( jamp_sv[ip * ncolor + icol] )[ieppV]; + JI[off + icol] = cximag( jamp_sv[ip * ncolor + icol] )[ieppV]; +#else + JR[off + icol] = cxreal( jamp_sv[ip * ncolor + icol] ); + JI[off + icol] = cximag( jamp_sv[ip * ncolor + icol] ); +#endif + } + } + } + // Ztemp[ncolor][ncol] = ColorMatrix[ncolor][ncolor] * Jamps[ncolor][ncol], real and imaginary parts apart + blas_symm( ncolor, ncol, cfsym.value, JR, ZR ); + blas_symm( ncolor, ncol, cfsym.value, JI, ZI ); + // |M|^2 for one (helicity, event) is the dot product of one column of Jamps with one column of Ztemp + for( int j = 0; j < ncol; j++ ) + { + const size_t off = (size_t)j * ncolor; + fptype2 me = 0; + for( int icol = 0; icol < ncolor; icol++ ) + me += JR[off + icol] * ZR[off + icol] + JI[off + icol] * ZI[off + icol]; + MEcol[j] = me; // may underflow #831 + } + // *** STORE THE RESULTS *** + // NB: MEs_ighel carries the running sum over helicities of |M|^2, which the + // event by event choice of helicity needs. The color sum is no longer added + // to allMEs one helicity at a time, so build those running sums here, + // starting from whatever allMEs already held (fix #435). + using E_ACCESS = HostAccessMatrixElements; // non-trivial access: buffer includes all events + for( int ip = 0; ip < nParityCS; ip++ ) + { + fptype_sv* running = ( ip == 0 ? MEs_ighel : MEs_ighel2 ); + fptype* MEsp = E_ACCESS::ieventAccessRecord( allMEs, ievt0 + ip * neppV ); + fptype_sv& MEsp_sv = E_ACCESS::kernelAccess( MEsp ); + for( int ieppV = 0; ieppV < neppV; ieppV++ ) + { +#ifdef MGONGPU_CPPSIMD + fptype sum = MEsp_sv[ieppV]; + for( int ighel = 0; ighel < nGoodHel; ighel++ ) + { + sum += MEcol[ighel * nevtB + ip * neppV + ieppV]; + running[ighel][ieppV] = sum; + } +#else + fptype sum = MEsp_sv; + for( int ighel = 0; ighel < nGoodHel; ighel++ ) + { + sum += MEcol[ighel * nevtB + ip * neppV + ieppV]; + running[ighel] = sum; + } +#endif + } + MEsp_sv = running[nGoodHel - 1]; + } + } +#endif +#endif diff --git a/madgraph/iolibs/template_files/madmatrix/color_sum_blas_loop.inc b/madgraph/iolibs/template_files/madmatrix/color_sum_blas_loop.inc new file mode 100644 index 0000000000..6de73a4f45 --- /dev/null +++ b/madgraph/iolibs/template_files/madmatrix/color_sum_blas_loop.inc @@ -0,0 +1,23 @@ +#ifdef MGONGPU_CPP_HAS_BLAS + // The color matrix does not depend on the helicity, so keep the jamps of + // every good helicity and hand them all to BLAS in a single call after + // the loop (see color_sum_cpu_blas). Everything else the loop was doing - + // the jamp2 sums for the color selection, the numerators and the + // denominators - stays inside calculate_jamps and still happens once per + // helicity, exactly as before. + static thread_local std::vector ghelJamp_sv( (size_t)ncomb * nParity * ncolor ); + for( int ighel = 0; ighel < cNGoodHel; ighel++ ) + { + const int ihel = cGoodHel[ighel]; + cxtype_sv* jamp_sv = ghelJamp_sv.data() + (size_t)ighel * nParity * ncolor; + for( int i = 0; i < nParity * ncolor; i++ ) jamp_sv[i] = cxzero_sv(); // calculate_jamps accumulates into jamp_sv + // **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 ); + } +#if defined MGONGPU_CPPSIMD and defined MGONGPU_FPTYPE_DOUBLE and defined MGONGPU_FPTYPE2_FLOAT + color_sum_cpu_blas( allMEs, MEs_ighel, MEs_ighel2, ghelJamp_sv.data(), cNGoodHel, ievt00 ); +#else + color_sum_cpu_blas( allMEs, MEs_ighel, nullptr, ghelJamp_sv.data(), cNGoodHel, ievt00 ); +#endif +#else diff --git a/madgraph/iolibs/template_files/madmatrix/madmatrix.mk b/madgraph/iolibs/template_files/madmatrix/madmatrix.mk index c82234690e..17b34c4c0a 100644 --- a/madgraph/iolibs/template_files/madmatrix/madmatrix.mk +++ b/madgraph/iolibs/template_files/madmatrix/madmatrix.mk @@ -643,6 +643,39 @@ GPUFLAGS += $(BLASCXXFLAGS) #------------------------------------------------------------------------------- +#=== Configure defaults and check if user-defined choices exist for CPPBLAS + +# HASBLAS above is about cuBLAS/hipBLAS, which only a GPU build can use. CPPBLAS +# is the separate question of whether the C++ color sum goes through a host BLAS: +# the color matrix does not depend on the helicity, so all the good helicities of +# one event page are the columns of one SYMM call. Whether a BLAS carrying SYMM +# could be linked was settled when this directory was written out; it is only +# taken for processes whose color matrix is large enough to be worth it, and for +# those the generated color_sum.cc carries both paths (example: "make CPPBLAS=hasNoBlas"). +ifeq ($(CPPBLAS),) + ifeq ($(GPUCC),) # CPU-only build + override CPPBLAS = %(cpp_blas_default)s + else # the GPU build does its color sum on the device + override CPPBLAS = hasNoBlas + endif +endif + +override CPPBLASCXXFLAGS= +override CPPBLASLIBFLAGS= + +ifeq ($(CPPBLAS),hasBlas) + override CPPBLASCXXFLAGS += -DMGONGPU_CPP_HAS_BLAS + override CPPBLASLIBFLAGS += %(cpp_blas_libflags)s +else ifneq ($(CPPBLAS),hasNoBlas) + $(error Unknown CPPBLAS='$(CPPBLAS)': only 'hasBlas' and 'hasNoBlas' are supported) +endif +CXXFLAGS += $(CPPBLASCXXFLAGS) + +#$(info CPPBLAS=$(CPPBLAS)) +#$(info CPPBLASLIBFLAGS=$(CPPBLASLIBFLAGS)) + +#------------------------------------------------------------------------------- + #=== Configure Position-Independent Code CXXFLAGS += -fPIC GPUFLAGS += $(XCOMPILERFLAG) -fPIC @@ -773,7 +806,7 @@ $(LIBDIR)/lib$(MADMATRIX_COMMONLIB).so: $(SRC)/*.h $(SRC)/*.cc $(BUILDDIR)/.buil # Target (and build rules): process shared library (C++ or CUDA/HIP, selected by GPUCC) ifeq ($(GPUCC),) $(LIBDIR)/lib$(MADMATRIX_LIB).so: $(LIBDIR)/lib$(MADMATRIX_COMMONLIB).so $(objects_lib) - $(CXX) -shared -o $@ $(objects_lib) $(CXXLIBFLAGSRPATH2) -L$(LIBDIR) -l$(MADMATRIX_COMMONLIB) + $(CXX) -shared -o $@ $(objects_lib) $(CXXLIBFLAGSRPATH2) -L$(LIBDIR) -l$(MADMATRIX_COMMONLIB) $(CPPBLASLIBFLAGS) else $(LIBDIR)/lib$(MADMATRIX_LIB).so: $(LIBDIR)/lib$(MADMATRIX_COMMONLIB).so $(objects_lib) $(GPUCC) --shared -o $@ $(objects_lib) $(GPULIBFLAGSRPATH2) -L$(LIBDIR) -l$(MADMATRIX_COMMONLIB) $(BLASLIBFLAGS) diff --git a/madgraph/iolibs/template_files/madmatrix/madmatrix_standalone.mk b/madgraph/iolibs/template_files/madmatrix/madmatrix_standalone.mk index 618cb9e42a..2d9c6a1525 100644 --- a/madgraph/iolibs/template_files/madmatrix/madmatrix_standalone.mk +++ b/madgraph/iolibs/template_files/madmatrix/madmatrix_standalone.mk @@ -35,7 +35,7 @@ standalone_all: all.$(TAG) check_sa.exe # code (the AOSOA->SoA transposition kernel). ifeq ($(GPUCC),) check_sa.exe: $(standalone_objects) $(LIBDIR)/lib$(MADMATRIX_LIB).so $(LIBDIR)/lib$(MADMATRIX_COMMONLIB).so - $(CXX) -o $@ $(standalone_objects) $(CXXLIBFLAGSRPATH) -L$(LIBDIR) -l$(MADMATRIX_LIB) -l$(MADMATRIX_COMMONLIB) $(BLASLIBFLAGS) + $(CXX) -o $@ $(standalone_objects) $(CXXLIBFLAGSRPATH) -L$(LIBDIR) -l$(MADMATRIX_LIB) -l$(MADMATRIX_COMMONLIB) $(BLASLIBFLAGS) $(CPPBLASLIBFLAGS) else check_sa.exe: $(standalone_objects) $(LIBDIR)/lib$(MADMATRIX_LIB).so $(LIBDIR)/lib$(MADMATRIX_COMMONLIB).so $(GPUCC) -o $@ $(standalone_objects) $(GPULIBFLAGSRPATH) -L$(LIBDIR) -l$(MADMATRIX_LIB) -l$(MADMATRIX_COMMONLIB) $(BLASLIBFLAGS) diff --git a/madgraph/iolibs/template_files/madmatrix/process_sigmaKin_function.inc b/madgraph/iolibs/template_files/madmatrix/process_sigmaKin_function.inc index 227301a6ea..7582ed1d99 100644 --- a/madgraph/iolibs/template_files/madmatrix/process_sigmaKin_function.inc +++ b/madgraph/iolibs/template_files/madmatrix/process_sigmaKin_function.inc @@ -116,7 +116,7 @@ #if defined MGONGPU_CPPSIMD and defined MGONGPU_FPTYPE_DOUBLE and defined MGONGPU_FPTYPE2_FLOAT fptype_sv MEs_ighel2[ncomb] = {}; // sum of MEs for all good helicities up to ighel (for the second neppV page) #endif - for( int ighel = 0; ighel < cNGoodHel; ighel++ ) +%(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) @@ -128,7 +128,7 @@ #if defined MGONGPU_CPPSIMD and defined MGONGPU_FPTYPE_DOUBLE and defined MGONGPU_FPTYPE2_FLOAT MEs_ighel2[ighel] = E_ACCESS::kernelAccess( E_ACCESS::ieventAccessRecord( allMEs, ievt00 + neppV ) ); #endif - } + }%(cpp_blas_helicity_loop_end)s // Event-by-event random choice of helicity #403 for( int ieppV = 0; ieppV < neppV; ++ieppV ) { diff --git a/madmatrix/model_handling.py b/madmatrix/model_handling.py index b4d1bac787..9d855da7e8 100644 --- a/madmatrix/model_handling.py +++ b/madmatrix/model_handling.py @@ -1483,6 +1483,8 @@ class OneProcessExporterMadMatrix(export_mg7.OneProcessExporterMG7): process_wavefunction_template = pjoin('madmatrix', 'cpp_process_wavefunctions.inc') process_sigmaKin_function_template = pjoin('madmatrix', 'process_sigmaKin_function.inc') single_process_template = pjoin('madmatrix', 'process_matrix.inc') + blas_color_sum_template = pjoin('madmatrix', 'color_sum_blas.inc') + blas_helicity_loop_template = pjoin('madmatrix', 'color_sum_blas_loop.inc') support_multichannel = False multichannel_var = ',fptype& multi_chanel_num, fptype& multi_chanel_denom' imaginary_unit = "cxtype(0,1)" @@ -1771,6 +1773,14 @@ def get_sigmaKin_lines(self, color_amplitudes, write=True): replace_dict['nb_channel'] = len(self.multi_channel_map) replace_dict['nb_color'] = max(1, len(self.matrix_elements[0].get('color_basis'))) + replace_dict['cpp_blas_helicity_loop'] = '' + replace_dict['cpp_blas_helicity_loop_end'] = '' + if self.cpp_blas_wanted(): + replace_dict['cpp_blas_helicity_loop'] = \ + self.read_template_file(self.blas_helicity_loop_template) + replace_dict['cpp_blas_helicity_loop_end'] = \ + '\n#endif // MGONGPU_CPP_HAS_BLAS' + if write: file = self.read_template_file(self.process_sigmaKin_function_template) % replace_dict file = strip_banner(file, banner_mark = "!") # skip first 8 lines in process_sigmaKin_function.inc (copyright) @@ -1984,6 +1994,24 @@ def edit_processidfile(self): ff.write(template % replace_dict) ff.close() + # AV - new method + @classmethod + def cpp_blas_wanted_for(cls, ncolor): + """Whether the C++ color sum goes through a host BLAS: only when one + carrying SYMM can be linked, and when the color matrix is big enough + that the call is worth setting up. Both the probe and the threshold are + the ones the Fortran color sum already uses. With BLAS off nothing is + written out, so color_sum.cc and CPPProcess.cc are character for + character the files written before any of this existed.""" + from madgraph.iolibs.export_v4 import ProcessExporterFortran + if not ProcessExporterFortran.blas_is_available(): + return False + return ncolor >= ProcessExporterFortran.blas_min_ncolor + + def cpp_blas_wanted(self): + return self.cpp_blas_wanted_for( + max(1, len(self.matrix_elements[0].get('color_basis')))) + # AV - new method def edit_colorsum(self): """Generate color_sum.cc""" @@ -1992,6 +2020,11 @@ def edit_colorsum(self): 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]) + replace_dict['cpp_blas_color_sum'] = '' + if self.cpp_blas_wanted(): + replace_dict['cpp_blas_color_sum'] = strip_banner( + open(pjoin(self.template_path, self.blas_color_sum_template), 'r').read(), + banner_mark='/') ff = open(pjoin(self.path, 'color_sum.cc'),'w') ff.write(template % replace_dict) ff.close() diff --git a/madmatrix/output.py b/madmatrix/output.py index 1820b10e96..74a33beb4d 100644 --- a/madmatrix/output.py +++ b/madmatrix/output.py @@ -159,6 +159,19 @@ def _parse_flavor_mask_option(self): return val.strip().lower() not in ('false', '0', 'no', 'off') return bool(val) + def get_makefile_replace_dict(self, model): + """Add what madmatrix.mk needs to know about a host BLAS for the C++ + color sum. Whether a given process actually takes it is decided when + that process is written out (see cpp_blas_wanted); this only settles + whether one could be linked at all.""" + + from madgraph.iolibs.export_v4 import ProcessExporterFortran + replace_dict = super().get_makefile_replace_dict(model) + flags = ProcessExporterFortran.blas_available_flags() + replace_dict['cpp_blas_default'] = 'hasBlas' if flags else 'hasNoBlas' + replace_dict['cpp_blas_libflags'] = flags + return replace_dict + # AV - overload the default version: create CMake directory, do not create lib directory def copy_template(self, model): misc.sprint('Entering ProcessExporterMadMatrix.copy_template (initialise the directory)') @@ -230,10 +243,8 @@ class ProcessExporterMadMatrixStandalone(ProcessExporterMadMatrix): def copy_template(self, model): super().copy_template(model) madmatrix_mk = pjoin(self.madmatrix_templates, 'madmatrix.mk') - rendered = self.read_template_file(madmatrix_mk) % { - 'model': self.get_model_name(model.get('name')), - 'cpp_compiler': self.opt['cpp_compiler'] if self.opt['cpp_compiler'] else 'g++', - } + rendered = self.read_template_file(madmatrix_mk) % \ + self.get_makefile_replace_dict(model) open(pjoin(self.dir_path, 'SubProcesses', 'madmatrix.mk'), 'w').write(rendered) # Write another custom bin/generate_events to orchestrate the standalone mode From 4bdf03957080a123156fd6229fb2a1e43be2ebac Mon Sep 17 00:00:00 2001 From: Olivier Mattelaer Date: Fri, 7 Aug 2026 23:57:21 +0200 Subject: [PATCH 39/42] fold the cuBLAS color sum onto one flow per reversal pair The previous commit folded color_sum_cpu and color_sum_kernel but left the BLAS color sum doing the full ncolor x ncolor GEMM: it multiplied the folded matrix spread back over the ncolor flows, with the dropped rows and columns at zero. Correct, and it could not drift away from the kernel, but roughly four times more arithmetic than needed. Folding it means compacting the jamps, which the representative flows being non-contiguous makes a gather and not a copy. convertD2F_Jamps becomes gatherFold_Jamps: it reads through colorFoldRep and is compiled in every precision mode, not just the mixed one, since the conversion to fptype2 is a no-op when the two types are the same. The gathered jamps go in a second buffer carved out of ghelAllBlasTmp, and all four BLAS calls then run over ncolorfold against s_pNormalizedColorMatrixFold2. With nothing left reading it, UnfoldedNormalizedColorMatrix and s_pNormalizedColorMatrix2 are dropped, and with them hostColorFoldRep, whose only consumer they were. ncolorfold lived only in color_sum.cc and the buffer sizing needs it elsewhere, so it is exported as CPPProcess::ncolorfold from process_class.inc and color_sum.cc reads it back from there, as it already did for ncolor. A process which does not fold (ncolorfold == ncolor) would otherwise pay for this twice, in a scratch buffer of double the size and in an identity gather which copies the jamps onto themselves. So the gather and its buffer are taken only when the gather is not the identity or there is a conversion to do; those processes run exactly the code they ran before. The rule and the buffer size are stated once, in blasColorSumTmpSize, which both the allocation in MatrixElementKernels.cc and the reset in color_sum_gpu now call instead of carrying a copy of the formula each. Scratch buffer, nGoodHel * nevt fptype2 per unit: g g > g g g 24 -> 12 flows: 48 -> 48 (d/f), 97 -> 49 (mixed) g g > g g g g 120 -> 60 flows: 240 -> 240 (d/f), 481 -> 241 (mixed) u u~ > u u~ g no folding: 8 -> 8 (d/f), 17 -> 17 (mixed) NB: NOT validated on a GPU - no CUDA toolkit was available, so nvcc has never seen these sources and neither has cuBLAS. What was done instead is to compile the GPU branch of the generated color_sum.cc with a host shim (empty __device__ and __global__, emulated gridDim/threadIdx, a reference column-major GEMM honouring the op/ld/stride semantics) and run color_sum_gpu down both of its paths on the same jamps. Over 3 processes x 3 precision modes the BLAS color sum agrees with the kernel one (which is what HASBLAS=hasNoBlas runs) to 6.1e-16 in double and 6.0e-07 in float, agrees with a long double reference to the same accuracy, and reproduces the unfolded BLAS color sum it replaces exactly. A canary past the end of the scratch buffer is untouched in all nine. Four negative controls turn that test red: a gather ignoring colorFoldRep, a batched GEMM stride left at ncolor, an under-sized buffer, and the no-gather shortcut taken on a process which does fold. The CPU color sum is untouched and ./check_sa.exe is bit-identical for g g > g g g. colorFoldRep itself has not been through nvcc either - it came in with the previous commit. Co-Authored-By: Claude Opus 5 --- .../madmatrix/MatrixElementKernels.cc | 13 +- .../template_files/madmatrix/color_sum.cc | 169 ++++++++---------- .../template_files/madmatrix/color_sum.h | 38 ++++ .../madmatrix/process_class.inc | 1 + madmatrix/model_handling.py | 47 +++-- 5 files changed, 155 insertions(+), 113 deletions(-) diff --git a/madgraph/iolibs/template_files/madmatrix/MatrixElementKernels.cc b/madgraph/iolibs/template_files/madmatrix/MatrixElementKernels.cc index 872e4795e3..207668c89a 100644 --- a/madgraph/iolibs/template_files/madmatrix/MatrixElementKernels.cc +++ b/madgraph/iolibs/template_files/madmatrix/MatrixElementKernels.cc @@ -10,6 +10,7 @@ #include "GpuRuntime.h" // Includes the abstraction for Nvidia/AMD compilation #include "MemoryAccessMomenta.h" #include "MemoryBuffers.h" +#include "color_sum.h" // for blasColorSumTmpSize #include // for fetestexcept #include @@ -464,14 +465,10 @@ namespace mg5amcGpu m_pHelNumerators.reset( new DeviceBufferSimple( nGoodHel * CPPProcess::ndiagrams * nevt ) ); m_pHelDenominators.reset( new DeviceBufferSimple( nGoodHel * nevt ) ); #ifndef MGONGPU_HAS_NO_BLAS - // Create the "many-helicity" super-buffers of real/imag ncolor*nevt temporary buffers for cuBLAS/hipBLAS intermediate results in color_sum_blas -#if defined MGONGPU_FPTYPE_DOUBLE and defined MGONGPU_FPTYPE2_FLOAT - // Mixed precision mode: need two fptype2[ncolor*2*nevt] buffers and one fptype2[nevt] buffer per good helicity - if( m_blasColorSum ) m_pHelBlasTmp.reset( new DeviceBufferSimple2( nGoodHel * ( 2 * CPPProcess::ncolor * mgOnGpu::nx2 + 1 ) * nevt ) ); -#else - // Standard single/double precision mode: need one fptype2[ncolor*2*nevt] buffer per good helicity - if( m_blasColorSum ) m_pHelBlasTmp.reset( new DeviceBufferSimple2( nGoodHel * CPPProcess::ncolor * mgOnGpu::nx2 * nevt ) ); -#endif + // Create the "many-helicity" super-buffer of temporary buffers for the cuBLAS/hipBLAS intermediate + // results in color_sum_blas, and for the jamps gathered onto the ncolorfold color flows the color + // sum is folded onto (see blasColorSumTmpSize in color_sum.h, which is where the size is defined) + if( m_blasColorSum ) m_pHelBlasTmp.reset( new DeviceBufferSimple2( blasColorSumTmpSize( nGoodHel, nevt ) ) ); #endif // Return the number of good helicities return nGoodHel; diff --git a/madgraph/iolibs/template_files/madmatrix/color_sum.cc b/madgraph/iolibs/template_files/madmatrix/color_sum.cc index 5340b9f5c0..a8cc2a9855 100644 --- a/madgraph/iolibs/template_files/madmatrix/color_sum.cc +++ b/madgraph/iolibs/template_files/madmatrix/color_sum.cc @@ -39,27 +39,6 @@ namespace mg5amcCpu }; // The fptype2 version is the default used by kernels (supporting mixed floating point mode) static __device__ fptype2 s_pNormalizedColorMatrixFold2[ncolorfold * ncolorfold]; -#ifndef MGONGPU_HAS_NO_BLAS - // The same matrix spread back over the ncolor unfolded color flows, which is what BLAS - // multiplies: the rows and columns of the flows which are not kept are left at zero, so - // the dropped flows contribute nothing and the product is the folded sum written out in - // full. The BLAS color sum is therefore NOT folded - it does the same ncolor x ncolor - // work it did before - but it takes its numbers from the one folded matrix which is - // written out, so it cannot drift away from what the kernel computes. - template - struct UnfoldedNormalizedColorMatrix - { - constexpr __host__ __device__ UnfoldedNormalizedColorMatrix() - : value() - { - for( int ifold = 0; ifold < ncolorfold; ifold++ ) - for( int jfold = 0; jfold < ncolorfold; jfold++ ) - value[hostColorFoldRep[ifold] * ncolor + hostColorFoldRep[jfold]] = colorMatrix[ifold][jfold] / colorDenom[ifold]; - } - T value[ncolor * ncolor]; - }; - static __device__ fptype2 s_pNormalizedColorMatrix2[ncolor * ncolor]; -#endif #endif //-------------------------------------------------------------------------- @@ -73,10 +52,6 @@ namespace mg5amcCpu first = false; constexpr NormalizedColorMatrix normalizedColorMatrix2; gpuMemcpyToSymbol( s_pNormalizedColorMatrixFold2, normalizedColorMatrix2.value, ncolorfold * ncolorfold * sizeof( fptype2 ) ); -#ifndef MGONGPU_HAS_NO_BLAS - constexpr UnfoldedNormalizedColorMatrix unfoldedNormalizedColorMatrix2; - gpuMemcpyToSymbol( s_pNormalizedColorMatrix2, unfoldedNormalizedColorMatrix2.value, ncolor * ncolor * sizeof( fptype2 ) ); -#endif } } #endif @@ -241,23 +216,45 @@ namespace mg5amcCpu #ifdef MGONGPUCPP_GPUIMPL #ifndef MGONGPU_HAS_NO_BLAS -#if defined MGONGPU_FPTYPE_DOUBLE and defined MGONGPU_FPTYPE2_FLOAT + // Compact the jamps onto the ncolorfold color flows the sum runs over, one per reversal + // pair when the color basis folds (see color_sum_blas): the representative flows are not + // contiguous, so this is a gather through colorFoldRep and not a copy. Without a folding + // ncolorfold == ncolor and colorFoldRep is the identity, so this only converts. The + // conversion to fptype2 is a no-op except in mixed floating point mode (double to float). __global__ void - convertD2F_Jamps( fptype2* allJampsFpt2, // output: jamp[2][ncolor][ihel][nevt] for one specific helicity ihel + gatherFold_Jamps( fptype2* allJampsFold2, // output: jamp[2][ncolorfold][ihel][nevt] for one specific helicity ihel const fptype* allJamps, // input: jamp[2][ncolor][ihel][nevt] for one specific helicity ihel const int nhel ) // input: number of good helicities nGoodHel { const int nevt = gridDim.x * blockDim.x; const int ievt = blockDim.x * blockIdx.x + threadIdx.x; constexpr int ihel = 0; // the input buffer allJamps already points to a specific helicity - // NB! From a functional point of view, any striding will be ok here as long as ncolor*2*nevt elements are all correctly copied! + // NB! From a functional point of view, any striding will be ok here as long as ncolorfold*2*nevt elements are all correctly gathered! // NB! Just in case this may be better for performance reasons, however, the same striding as in compute_jamps and cuBLAS is used here for( int ix2 = 0; ix2 < mgOnGpu::nx2; ix2++ ) - for( int icol = 0; icol < ncolor; icol++ ) - allJampsFpt2[ix2 * ncolor * nhel * nevt + icol * nhel * nevt + ihel * nevt + ievt] = - allJamps[ix2 * ncolor * nhel * nevt + icol * nhel * nevt + ihel * nevt + ievt]; + for( int ifold = 0; ifold < ncolorfold; ifold++ ) + allJampsFold2[ix2 * ncolorfold * nhel * nevt + ifold * nhel * nevt + ihel * nevt + ievt] = + allJamps[ix2 * ncolor * nhel * nevt + colorFoldRep[ifold] * nhel * nevt + ihel * nevt + ievt]; + } + + // Gather the jamps of every good helicity into ghelAllJampsBuf, and return it + fptype2* + gatherFold_AllJamps( fptype2* ghelAllJampsBuf, // output: allJamps super-buffer[2][ncolorfold][nhel][nevt] + const fptype* ghelAllJamps, // input: allJamps super-buffer[2][ncolor][nhel][nevt] + gpuStream_t* ghelStreams, // input: cuda streams (index is ighel) + const int nhel, // input: number of good helicities + const int gpublocks, // input: cuda gpublocks + const int gputhreads ) // input: cuda gputhreads + { + const int nevt = gpublocks * gputhreads; + for( int ighel = 0; ighel < nhel; ighel++ ) + { + const fptype* hAllJamps = ghelAllJamps + ighel * nevt; // jamps for a single helicity ihel + fptype2* hAllJampsFold2 = ghelAllJampsBuf + ighel * nevt; // folded jamps for a single helicity ihel + gpuLaunchKernelStream( gatherFold_Jamps, gpublocks, gputhreads, ghelStreams[ighel], hAllJampsFold2, hAllJamps, nhel ); + } + return ghelAllJampsBuf; } -#endif #endif #endif @@ -286,67 +283,61 @@ namespace mg5amcCpu const fptype* ghelAllJamps, // input: allJamps super-buffer[2][ncol][nhel][nevt] for nhel good helicities fptype2* ghelAllBlasTmp, // tmp: allBlasTmp super-buffer for nhel good helicities gpuBlasHandle_t* pBlasHandle, // input: cuBLAS/hipBLAS handle -#if defined MGONGPU_FPTYPE_DOUBLE and defined MGONGPU_FPTYPE2_FLOAT gpuStream_t* ghelStreams, // input: cuda streams (index is ighel: only the first nhel <= ncomb are non-null) -#else - gpuStream_t* /*ghelStreams*/, // input: cuda streams (index is ighel: only the first nhel <= ncomb are non-null) -#endif const int nhel, // input: number of good helicities (nhel == nGoodHel) const int gpublocks, // input: cuda gpublocks const int gputhreads ) // input: cuda gputhreads { const int nevt = gpublocks * gputhreads; - // NB: unlike color_sum_cpu and color_sum_kernel, the BLAS color sum is NOT folded onto one - // color flow per reversal pair. Folding it would mean compacting allJamps from ncolor down - // to ncolorfold, which the jamps are not written in, so it would take a gather kernel and a - // second buffer. Instead the matrix it multiplies is the folded one spread back over the - // ncolor flows (see UnfoldedNormalizedColorMatrix): same numbers, same ncolor x ncolor work - // as before. Compacting the jamps would make this a factor 4 cheaper and is left to do. + // As in color_sum_cpu and color_sum_kernel, the sum is folded onto one color flow per + // reversal pair: the jamps are first gathered from the ncolor flows they are written in + // down to the ncolorfold flows which are kept (see gatherFold_Jamps), and it is those + // which are multiplied by the folded color matrix. Without a folding ncolorfold == ncolor, + // the gather is the identity and this is the plain ncolor x ncolor color sum. - // Get the address associated with the normalized color matrix in device memory + // Get the address associated with the normalized folded color matrix in device memory static fptype2* devNormColMat = nullptr; - if( !devNormColMat ) gpuGetSymbolAddress( (void**)&devNormColMat, s_pNormalizedColorMatrix2 ); + if( !devNormColMat ) gpuGetSymbolAddress( (void**)&devNormColMat, s_pNormalizedColorMatrixFold2 ); + // The scratch buffer holds the BLAS intermediate results, the gathered jamps if they need + // a buffer of their own, and in mixed precision mode the fptype2 MEs: see the layout in + // blasColorSumTmpSize, which is what MatrixElementKernels.cc allocates. + fptype2* ghelAllZtempBoth = ghelAllBlasTmp; // start of the fptype2[ncolorfold*2*nhel*nevt] buffer + fptype2* ghelAllJampsBuf = ghelAllBlasTmp + ncolorfold * mgOnGpu::nx2 * nhel * nevt; // start of the second one, if there is one #if defined MGONGPU_FPTYPE_DOUBLE and defined MGONGPU_FPTYPE2_FLOAT - // Mixed precision mode: need two fptype2[2*ncolor*nhel*nevt] buffers and one fptype2[nhel*nevt] buffers for the nhel helicities - fptype2* ghelAllZtempBoth = ghelAllBlasTmp; // start of first fptype2[ncolor*2*nhel*nevt] buffer - fptype2* ghelAllJampsFpt2 = ghelAllBlasTmp + ncolor * mgOnGpu::nx2 * nhel * nevt; // start of second fptype2[ncolor*2*nhel*nevt] buffer - fptype2* ghelAllMEsFpt2 = ghelAllBlasTmp + 2 * ncolor * mgOnGpu::nx2 * nhel * nevt; // start of fptype2[nhel*nevt] buffer - // Convert jamps from double to float - for( int ighel = 0; ighel < nhel; ighel++ ) - { - const fptype* hAllJamps = ghelAllJamps + ighel * nevt; // jamps for a single helicity ihel - fptype2* hAllJampsFpt2 = ghelAllJampsFpt2 + ighel * nevt; // jamps for a single helicity ihel - gpuLaunchKernelStream( convertD2F_Jamps, gpublocks, gputhreads, ghelStreams[ighel], hAllJampsFpt2, hAllJamps, nhel ); - } - // Real and imaginary components - const fptype2* ghelAllJampsReal = ghelAllJampsFpt2; - const fptype2* ghelAllJampsImag = ghelAllJampsFpt2 + ncolor * nhel * nevt; + // Mixed precision mode: the jamps are gathered into a buffer of their own in any case, + // as they must be converted from double to float on the way + static_assert( blasColorSumNeedsJampBuffer() ); + fptype2* ghelAllMEsFpt2 = ghelAllBlasTmp + 2 * ncolorfold * mgOnGpu::nx2 * nhel * nevt; // start of the fptype2[nhel*nevt] buffer + const fptype2* ghelAllJampsFold2 = gatherFold_AllJamps( ghelAllJampsBuf, ghelAllJamps, ghelStreams, nhel, gpublocks, gputhreads ); #else - // Standard single or double precision mode: need one fptype2[ncolor*2*nhel*nevt] buffer static_assert( std::is_same::value ); - fptype2* ghelAllZtempBoth = ghelAllBlasTmp; // start of fptype2[ncolor*2*nhel*nevt] buffer fptype2* ghelAllMEsFpt2 = ghelAllMEs; - // Real and imaginary components - const fptype2* ghelAllJampsReal = ghelAllJamps; // this is not a cast (the two types are identical) - const fptype2* ghelAllJampsImag = ghelAllJamps + ncolor * nhel * nevt; // this is not a cast (the two types are identical) + // Without a folding there is nothing to gather and nothing to convert: read the jamps + // where compute_jamps wrote them (this is not a cast, the two types are identical) + const fptype2* ghelAllJampsFold2 = + ( blasColorSumNeedsJampBuffer() + ? gatherFold_AllJamps( ghelAllJampsBuf, ghelAllJamps, ghelStreams, nhel, gpublocks, gputhreads ) + : ghelAllJamps ); #endif // Real and imaginary components + const fptype2* ghelAllJampsReal = ghelAllJampsFold2; + const fptype2* ghelAllJampsImag = ghelAllJampsFold2 + ncolorfold * nhel * nevt; fptype2* ghelAllZtempReal = ghelAllZtempBoth; - fptype2* ghelAllZtempImag = ghelAllZtempBoth + ncolor * nhel * nevt; + fptype2* ghelAllZtempImag = ghelAllZtempBoth + ncolorfold * nhel * nevt; - // Note: striding for cuBLAS from DeviceAccessJamp: - // - ghelAllJamps(icol,ihel,ievt).real is ghelAllJamps[0 * ncolor * nhel * nevt + icol * nhel * nevt + ihel * nevt + ievt] - // - ghelAllJamps(icol,ihel,ievt).imag is ghelAllJamps[1 * ncolor * nhel * nevt + icol * nhel * nevt + ihel * nevt + ievt] + // Note: striding for cuBLAS from gatherFold_Jamps (that of DeviceAccessJamp, over ncolorfold): + // - ghelAllJampsFold2(ifold,ihel,ievt).real is ghelAllJampsFold2[0 * ncolorfold * nhel * nevt + ifold * nhel * nevt + ihel * nevt + ievt] + // - ghelAllJampsFold2(ifold,ihel,ievt).imag is ghelAllJampsFold2[1 * ncolorfold * nhel * nevt + ifold * nhel * nevt + ihel * nevt + ievt] - // Step 1: Compute Ztemp[ncolor][nhel*nevt] = ColorMatrix[ncolor][ncolor] * JampsVector[ncolor][nhel*nevt] for both real and imag + // Step 1: Compute Ztemp[ncolorfold][nhel*nevt] = ColorMatrix[ncolorfold][ncolorfold] * JampsVector[ncolorfold][nhel*nevt] for both real and imag // In this case alpha=1 and beta=0: the operation is Ztemp = alpha * ColorMatrix * JampsVector + beta * Ztemp fptype2 alpha1 = 1; fptype2 beta1 = 0; - const int ncolorM = ncolor; + const int ncolorM = ncolorfold; const int nevtN = nhel*nevt; - const int ncolorK = ncolor; + const int ncolorK = ncolorfold; checkGpuBlas( gpuBlasTgemm( *pBlasHandle, GPUBLAS_OP_N, // do not transpose ColMat GPUBLAS_OP_T, // transpose JampsV (new1) @@ -366,31 +357,31 @@ namespace mg5amcCpu &beta1, ghelAllZtempImag, ncolorM ) ); // Ztemp is ncolorM x nevtN - // Step 2: For each ievt, compute the dot product of JampsVector[ncolor][ievt] dot tmp[ncolor][ievt] + // Step 2: For each ievt, compute the dot product of JampsVector[ncolorfold][ievt] dot tmp[ncolorfold][ievt] // In this case alpha=1 and beta=1: the operation is ME = alpha * ( Tmp dot JampsVector ) + beta * ME // Use cublasSgemmStridedBatched to perform these batched dot products in one call fptype2 alpha2 = 1; fptype2 beta2 = 1; checkGpuBlas( gpuBlasTgemmStridedBatched( *pBlasHandle, - GPUBLAS_OP_N, // do not transpose JampsV (new1) - GPUBLAS_OP_N, // do not transpose Tmp - 1, 1, ncolor, // result is 1x1 (dot product) + GPUBLAS_OP_N, // do not transpose JampsV (new1) + GPUBLAS_OP_N, // do not transpose Tmp + 1, 1, ncolorfold, // result is 1x1 (dot product) &alpha2, - ghelAllJampsReal, nevtN, 1, // allJamps is nevtN x ncolor, stride 1 for each ievt column - ghelAllZtempReal, ncolor, ncolor, // allZtemp is ncolor x nevtN, with stride ncolor for each ievt column + ghelAllJampsReal, nevtN, 1, // allJamps is nevtN x ncolorfold, stride 1 for each ievt column + ghelAllZtempReal, ncolorfold, ncolorfold, // allZtemp is ncolorfold x nevtN, with stride ncolorfold for each ievt column &beta2, - ghelAllMEsFpt2, 1, 1, // output is a 1x1 result for each "batch" (i.e. for each ievt) - nevtN ) ); // there are nevtN (nhel*nevt) "batches" + ghelAllMEsFpt2, 1, 1, // output is a 1x1 result for each "batch" (i.e. for each ievt) + nevtN ) ); // there are nevtN (nhel*nevt) "batches" checkGpuBlas( gpuBlasTgemmStridedBatched( *pBlasHandle, - GPUBLAS_OP_N, // do not transpose JampsV (new1) - GPUBLAS_OP_N, // do not transpose Tmp - 1, 1, ncolor, // result is 1x1 (dot product) + GPUBLAS_OP_N, // do not transpose JampsV (new1) + GPUBLAS_OP_N, // do not transpose Tmp + 1, 1, ncolorfold, // result is 1x1 (dot product) &alpha2, - ghelAllJampsImag, nevtN, 1, // allJamps is nevtN x ncolor, stride 1 for each ievt column (new1) - ghelAllZtempImag, ncolor, ncolor, // allZtemp is ncolor x nevtN, with stride ncolor for each ievt column + ghelAllJampsImag, nevtN, 1, // allJamps is nevtN x ncolorfold, stride 1 for each ievt column (new1) + ghelAllZtempImag, ncolorfold, ncolorfold, // allZtemp is ncolorfold x nevtN, with stride ncolorfold for each ievt column &beta2, - ghelAllMEsFpt2, 1, 1, // output is a 1x1 result for each "batch" (i.e. for each ievt) - nevtN ) ); // there are nevt (nhel*nevt) "batches" + ghelAllMEsFpt2, 1, 1, // output is a 1x1 result for each "batch" (i.e. for each ievt) + nevtN ) ); // there are nevt (nhel*nevt) "batches" #if defined MGONGPU_FPTYPE_DOUBLE and defined MGONGPU_FPTYPE2_FLOAT // Convert MEs from float to double @@ -447,12 +438,8 @@ namespace mg5amcCpu assert( false ); // BLAS in async mode not supported for now } else { checkGpu( gpuDeviceSynchronize() ); // do not start the BLAS color sum for all helicities until the loop over helicities has completed - // Reset the tmp buffer -#if defined MGONGPU_FPTYPE_DOUBLE and defined MGONGPU_FPTYPE2_FLOAT - gpuMemset( ghelAllBlasTmp, 0, nGoodHel * nevt * ( 2 * ncolor * mgOnGpu::nx2 + 1 ) * sizeof( fptype2 ) ); -#else - gpuMemset( ghelAllBlasTmp, 0, nGoodHel * nevt * ( ncolor * mgOnGpu::nx2 ) * sizeof( fptype2 ) ); -#endif + // Reset the tmp buffer (same size as the one MatrixElementKernelDevice allocated) + gpuMemset( ghelAllBlasTmp, 0, blasColorSumTmpSize( nGoodHel, nevt ) * sizeof( fptype2 ) ); // Delegate the color sum to BLAS for color_sum_blas( ghelAllMEs, ghelAllJamps, ghelAllBlasTmp, pBlasHandle, ghelStreams, nGoodHel, gpublocks, gputhreads ); } diff --git a/madgraph/iolibs/template_files/madmatrix/color_sum.h b/madgraph/iolibs/template_files/madmatrix/color_sum.h index 347184c4e1..8709554eeb 100644 --- a/madgraph/iolibs/template_files/madmatrix/color_sum.h +++ b/madgraph/iolibs/template_files/madmatrix/color_sum.h @@ -14,6 +14,8 @@ #include "CPPProcess.h" #include "GpuAbstraction.h" +#include + #ifdef MGONGPUCPP_GPUIMPL namespace mg5amcGpu #else @@ -22,6 +24,42 @@ namespace mg5amcCpu { //-------------------------------------------------------------------------- +#ifdef MGONGPUCPP_GPUIMPL +#ifndef MGONGPU_HAS_NO_BLAS + // The BLAS color sum multiplies the jamps gathered onto the ncolorfold color flows the sum + // runs over (see color_sum_blas): does that gather need a buffer of its own? Not when it is + // the identity because the color basis does not fold, and there is no fptype2 conversion to + // do either - there the jamps are read where compute_jamps already wrote them, exactly as + // the color sum did before it was folded. + constexpr bool + blasColorSumNeedsJampBuffer() + { +#if defined MGONGPU_FPTYPE_DOUBLE and defined MGONGPU_FPTYPE2_FLOAT + return true; // mixed precision mode: the jamps must be converted from double to float +#else + return CPPProcess::ncolorfold < CPPProcess::ncolor; +#endif + } + + // The size of the ghelAllBlasTmp scratch buffer color_sum_blas needs, in fptype2 elements: + // one fptype2[ncolorfold*nx2*nhel*nevt] buffer for the BLAS intermediate results, one more + // for the gathered jamps if they need one, and in mixed precision mode one fptype2[nhel*nevt] + // buffer for the MEs, which are fptype elsewhere. This is the one place the size is defined: + // both the allocation (MatrixElementKernels.cc) and the reset (color_sum_gpu) come here. + constexpr std::size_t + blasColorSumTmpSize( const int nhel, const int nevt ) + { + std::size_t nfptype2PerEvent = ( blasColorSumNeedsJampBuffer() ? 2 : 1 ) * CPPProcess::ncolorfold * mgOnGpu::nx2; +#if defined MGONGPU_FPTYPE_DOUBLE and defined MGONGPU_FPTYPE2_FLOAT + nfptype2PerEvent += 1; // the fptype2 matrix elements +#endif + return nfptype2PerEvent * (std::size_t)nhel * (std::size_t)nevt; + } +#endif +#endif + + //-------------------------------------------------------------------------- + #ifdef MGONGPUCPP_GPUIMPL class DeviceAccessJamp { diff --git a/madgraph/iolibs/template_files/madmatrix/process_class.inc b/madgraph/iolibs/template_files/madmatrix/process_class.inc index 59a6d07333..12720ae8f6 100644 --- a/madgraph/iolibs/template_files/madmatrix/process_class.inc +++ b/madgraph/iolibs/template_files/madmatrix/process_class.inc @@ -56,6 +56,7 @@ static constexpr int ncomb = %(nbhel)d; // #helicity combinations: e.g. 16 for e+ e- -> mu+ mu- (2**4 = fermion spin up/down ** npar) static constexpr int ndiagrams = %(ndiagrams)d; // #Feynman diagrams: e.g. 3 for e+ e- -> mu+ mu- static constexpr int ncolor = %(ncolor)s; // the number of leading colors: e.g. 1 for e+ e- -> mu+ mu- + static constexpr int ncolorfold = %(ncolorfold)s; // the number of color flows |M|^2 is summed over: one per reversal pair where the color basis folds, ncolor otherwise (see color_sum.cc) static constexpr int nmaxflavor = %(nmaxflavor)d; // the maximum number of flavor combinations // Hardcoded parameters for this process (constant class variables) diff --git a/madmatrix/model_handling.py b/madmatrix/model_handling.py index ae0b249fa2..114c31e74f 100644 --- a/madmatrix/model_handling.py +++ b/madmatrix/model_handling.py @@ -1516,6 +1516,7 @@ 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 + replace_dict['ncolorfold'] = self.get_ncolorfold(self.matrix_elements[0], replace_dict['ncolor']) replace_dict['nwave'] = 4 if (fd_gauge): replace_dict['nwave'] += 1 @@ -2137,6 +2138,29 @@ def write_process_cc_file(self, writer): def jamp_fold_worthwhile(self, sign, nb_pairs): return True + # AV - cache the export_v4.ColorReflectionFolding method + def get_jamp_folding(self, matrix_element): + """Cache the folding: it is read once for CPPProcess.h (ncolorfold) and + once for color_sum.cc, and finding it walks the whole color basis.""" + cache = self.__dict__.setdefault('_jamp_folding_cache', {}) + key = id(matrix_element) + if key not in cache: + # keep the matrix element alive so that its id cannot be reused + cache[key] = (matrix_element, + super().get_jamp_folding(matrix_element)) + return cache[key][1] + + def get_ncolorfold(self, matrix_element, ncolor): + """The number of color flows |M|^2 is summed over: one per reversal pair + where the color basis folds, every flow otherwise. Mirrors what + get_color_matrix_lines writes the folded color matrix over, and is + exported as CPPProcess::ncolorfold because the BLAS color sum sizes its + buffers on it outside color_sum.cc (see MatrixElementKernels.cc).""" + if not matrix_element.get('color_matrix'): + return 1 + folding = self.get_jamp_folding(matrix_element) + return len(folding['representatives']) if folding else ncolor + # AV - replace the export_cpp.OneProcessExporterCPP method (fix fptype and improve formatting) def get_color_matrix_lines(self, matrix_element): """Return the color matrix definition lines for this matrix element. Split rows in chunks of size n. @@ -2187,8 +2211,9 @@ def get_color_matrix_lines(self, matrix_element): @staticmethod def get_color_fold_lines(folding, ncolor): - """The number of color flows the sum runs over and which flow it keeps - out of every reversal pair. Without a folding this is every flow.""" + """Which color flow the sum keeps out of every reversal pair. Without a + folding this is every flow. How many there are is CPPProcess::ncolorfold + (see get_ncolorfold), which is where color_sum.cc reads it from.""" if folding: representatives = folding['representatives'] @@ -2206,22 +2231,16 @@ def get_color_fold_lines(folding, ncolor): chunks = [', '.join('%i' % line for line in representatives[start:start + 20]) for start in range(0, len(representatives), 20)] values = '{\n ' + ',\n '.join(chunks) + ' }' - # colorFoldRep is indexed at run time inside the GPU kernel, so it has to - # live in device memory, and a host copy is needed next to it: same split - # as channel2iconfig/hostChannel2iconfig in coloramps.h + # colorFoldRep is indexed at run time inside the GPU kernels, so it has + # to live in device memory: same split as channel2iconfig in coloramps.h + # (nvcc cannot read a plain constexpr array from device code without + # --expt-relaxed-constexpr, which the makefile does not pass) return comment + \ - ' constexpr int ncolorfold = %i; // the number of color flows |M|^2 is summed over\n' % len(representatives) + \ + ' constexpr int ncolorfold = CPPProcess::ncolorfold; // the number of color flows |M|^2 is summed over (%i here)\n' % len(representatives) + \ ' // Which color flow of each reversal pair is kept (C indexing, in [0, ncolor-1])\n' + \ ' // (NB: this array is created on the host in C++ code and on the device in GPU code)\n' + \ ' __device__ constexpr int colorFoldRep[ncolorfold] = %s; // 1-D array[%i]\n' \ - % (values, len(representatives)) + \ - '#ifdef MGONGPUCPP_GPUIMPL\n' + \ - ' // Host copy of the colorFoldRep array (needed to fold the color matrix at compile time)\n' + \ - ' constexpr int hostColorFoldRep[ncolorfold] = %s; // 1-D array[%i]\n' \ - % (values, len(representatives)) + \ - '#else\n' + \ - ' constexpr const int* hostColorFoldRep = colorFoldRep;\n' + \ - '#endif' + % (values, len(representatives)) # AV - replace the export_cpp.OneProcessExporterCPP method (improve formatting) def get_initProc_lines(self, matrix_element, color_amplitudes): From d718ce882050aa6822ef5e31149230f9099a3f09 Mon Sep 17 00:00:00 2001 From: Olivier Mattelaer Date: Sat, 8 Aug 2026 01:17:39 +0200 Subject: [PATCH 40/42] share the orbit equivariant color-flow search with the madmatrix backend The fortran exporters look for the shared JAMP sub-expressions by whole orbits of the permutations leaving the color basis invariant, not one at a time (jamp_orbit). The C++/cudacpp writer had none of it: it ran the plain greedy scan the previous commit gave it and stopped there. Everything about that search which is not printing now lives in jamp_optimiser.JampOptimiser next to the plain scan -- get_jamp_symmetry, optimise_jamp_equivariant, optimise_jamp_best and their helpers, about 360 lines. Two hooks are left for the backends, because it is the emission which decides what is usable: jamp_orbit_allowed (now also reading --jamp_orbit at output time) and jamp_greedy_tail_enabled, which stays 'jamp_emit == tables' for fortran since INIT_JAMP cannot rebuild sub-expressions that are not orbits of anything. What prints the result -- jamp_orbit_recipes, jamp_orbit_tables, get_jamp_decl_lines, get_jamp_init_routine, jamp_gather, namp_dim -- stays in export_v4. The fortran output is unchanged, byte for byte. optimise_jamp_best earns its keep straight away: on g g > g g g g over the DDM basis the plain scan is the shorter of the two (392 definitions against 399) and is kept, so that build comes out identical. Definitions and CPPProcess.cc, with the search off and on: g g > g g g g ddm 392 -> 392 250 220 -> 250 220 B g g > g g g g trace 951 -> 795 276 077 -> 266 886 B u u~ > u u~ ggg trace 756 -> 732 487 327 -> 482 790 B g g > t t~ ggg trace 3030 -> 2535 950 946 -> 899 241 B g g > g g g g g ddm 8524 -> 7560 2 392 147 -> 2 289 786 B The file moves less than the definitions do because the color flows are only 13 to 28% of CPPProcess.cc here, against 78% of the fortran matrix.f: the rest is the HELAS calls. |M|^2 is unchanged to the last digit on every process tested, and agrees with the fortran ./check at the same phase-space point to 6e-15 or better -- the pre-existing gap between the two backends, not something this adds. Throughput is unchanged within the noise of the machine; the color flows are not what these processes spend their time on. No table emission. The stage which could be driven by a table is the one building definitions out of other definitions, plus the color flows built out of them: both sides are arrays that already exist. The amplitudes are not -- each one passes through the single slot amp_sv[0] and is gone by the next diagram, so the capture stage is at least one written statement per amplitude whatever the search does, and giving it an array to index would cost nb_amp * cxtype_sv per SIMD page and per GPU thread (232 kB for g g > g g g g g). The table-drivable part peaks at 4515 entries on the largest process generated here, and a microbenchmark on the real cxtype_sv puts the crossover at about ten thousand -- twice the fortran figure of five thousand: 200 defs lines 0.071 us table 0.128 us ratio 1.81 1000 defs lines 0.355 us table 0.645 us ratio 1.82 4515 defs lines 2.329 us table 3.549 us ratio 1.52 8000 defs lines 7.678 us table 8.213 us ratio 1.07 16000 defs lines 18.74 us table 18.34 us ratio 0.98 32000 defs lines 48.14 us table 37.08 us ratio 0.77 So tables would be 1.5x slower on the biggest thing the backend has, for 8% of the source. jamp_emit stays fortran-only. --jamp_orbit=False recovers the plain scan, on either backend. Co-Authored-By: Claude Opus 5 --- madgraph/interface/madgraph_interface.py | 3 +- madgraph/iolibs/export_v4.py | 386 +------------------- madgraph/iolibs/jamp_optimiser.py | 428 ++++++++++++++++++++++- madmatrix/model_handling.py | 11 +- 4 files changed, 444 insertions(+), 384 deletions(-) diff --git a/madgraph/interface/madgraph_interface.py b/madgraph/interface/madgraph_interface.py index d07744e32c..9de0da38aa 100755 --- a/madgraph/interface/madgraph_interface.py +++ b/madgraph/interface/madgraph_interface.py @@ -523,6 +523,7 @@ def help_output(self): logger.info(" --noeps=True: no jpeg and eps diagrams will be generated.") logger.info(" -name: the postfix of the main file in pythia8 mode.") logger.info(" --jamp_optim=[True|False]: [madevent(default:True)|standalone(default:False)] allows a more efficient code computing the color-factor.") + logger.info(" --jamp_orbit=[True|False]: [madevent|standalone|mg7] look for the shared color-factor sub-expressions by whole orbits of the color basis symmetry.") logger.info(" --t_strategy: [madevent] allows to change ordering strategy for t-channel.") logger.info(" --hel_recycling=False: [madevent] forbids helicity recycling optimization") logger.info(" --mask=False: [madevent|standalone] disable flavor-mask optimization for grouped/merged flavors (default:True).") @@ -2694,7 +2695,7 @@ def complete_open(self, text, line, begidx, endidx): def complete_output(self, text, line, begidx, endidx, possible_options = ['f', 'noclean', 'nojpeg'], possible_options_full = ['-f', '-noclean', '-nojpeg', '--noeps=True','--hel_recycling=False', - '--jamp_optim=', '--t_strategy=', '--vector_size=4', '--nb_warp=1', + '--jamp_optim=', '--jamp_orbit=', '--t_strategy=', '--vector_size=4', '--nb_warp=1', '--mask=False', '--prefix=']): "Complete the output command" diff --git a/madgraph/iolibs/export_v4.py b/madgraph/iolibs/export_v4.py index 5945f0ad8d..75fdea3e19 100755 --- a/madgraph/iolibs/export_v4.py +++ b/madgraph/iolibs/export_v4.py @@ -407,31 +407,25 @@ class ProcessExporterFortran(ColorReflectionFolding, VirtualExporter, # jamp_fold (sum |M|^2 over one color flow per reversal pair) comes from # ColorReflectionFolding and stays off unless the template sums over # NCOLORFOLD: get_color_data_lines is shared by every fortran exporter. - # jamp_optim, myjamp_count and jamp_integer_walk come from JampOptimiser. + # jamp_optim, myjamp_count and jamp_integer_walk come from JampOptimiser, + # and so does the orbit equivariant optimisation itself: jamp_orbit, + # jamp_greedy_tail and jamp_compare_max_size. What is left here is how the + # definitions it produces are written out. # BLAS-3 for the color sum: all helicities at once as one right hand side. # None means take it when the library is there and the process is big # enough for it to pay. blas = None blas_min_ncolor = 100 - # write the JAMP definitions as one recipe per orbit of the permutations - # leaving the color basis invariant, instead of one line per definition - jamp_orbit = False # How the definitions reach memory: 'recipes' rebuilds them at the first # call from one recipe per orbit, 'tables' writes the operand indices out # as DATA. Both run the very same loop, and both start from the orbit # equivariant optimisation, so they only differ in the source they need. jamp_emit = 'tables' - # finish with the plain scan once the orbit rounds have nothing left to - # take as a whole (only used by the table emission, see below) - jamp_greedy_tail = True # Read the amplitudes of the current helicity into a buffer before running # the definitions over it, instead of holding the definitions at the end of # AMP. Needed where AMP is indexed by helicity, which is what madevent does # once it rewrites the matrix element for helicity recycling. jamp_gather = False - # up to this many entries in the matrix, both optimisations are run and the - # shorter result kept (see optimise_jamp_best) - jamp_compare_max_size = 20000 # Below this many definitions writing them out is both smaller and faster: # the lines still fit in the instruction cache, while the loop reading the # operands from a table pays for the two indirections whatever the size. @@ -3204,43 +3198,6 @@ def format(frac): return res_list, len(defs) - def optimise_jamp_best(self, all_element, symmetry): - """Taking whole orbits only pays once there is enough of them to share: - on a small matrix it can end up asking for more additions than the plain - scan, which is free to take whatever it likes. g g > t t~ g is such a - case, 46 additions against 39. - - Small matrices are cheap to optimise, so rather than guess where the - turn is, do both and keep the shorter. Above that size only the orbit - version is run: it wins by a wide margin on everything that big, and - the plain scan is the slow one there.""" - - orbit_element, orbit_defs = self.optimise_jamp_equivariant( - dict(all_element), symmetry) - if len(all_element) > self.jamp_compare_max_size: - return orbit_element, orbit_defs - - orbits = self.jamp_orbits - plain_element, plain_defs = self.optimise_jamp(dict(all_element)) - if self.jamp_operation_count(plain_element, plain_defs) < \ - self.jamp_operation_count(orbit_element, orbit_defs): - self.jamp_orbits = None - return plain_element, plain_defs - self.jamp_orbits = orbits - return orbit_element, orbit_defs - - #=========================================================================== - # Orbit equivariant version of the JAMP optimisation - #=========================================================================== - # A permutation of the external color indices which maps the color basis - # onto itself (see color_amp.ColorBasisSymmetry) also permutes the columns - # of the JAMP matrix, up to a sign. The whole matrix is then invariant, so - # the sub-expressions the optimisation looks for come in orbits: every one - # of them is worth exactly as much as the others. Introducing a whole orbit - # at a time, rather than one sub-expression at a time as the plain scan - # does, leaves the matrix invariant at every step, and the definitions can - # be written as one recipe per orbit. - _blas_available = None @classmethod @@ -3388,280 +3345,6 @@ def get_color_fold_ampso(self, folding, ncolor): 'color_fold_gather': " JFOLD(:,:) = JAMP(COLREP(:),:)", 'color_fold_array': 'JFOLD'} - @staticmethod - def jamp_column_form(column): - """Canonical form of one column of the JAMP matrix up to a global sign, - together with the sign which was taken out.""" - - entries = sorted(column.items()) - first = entries[0][1] - sign = -1 if (first.real, first.imag) < (0., 0.) else 1 - return tuple((i, sign * value) for i, value in entries), sign - - @classmethod - def jamp_amp_permutation(cls, columns, induced): - """Permutation of the amplitudes induced by the permutation induced of - the color basis: return {amp: (amp, sign)} such that - - M[induced[i], sigma(j)] = sign(j) * M[i, j] - - or None if the columns are not mapped onto each other. - - Several amplitudes often have the very same column, so the columns are - gathered by their canonical form and one target is taken out of each - group at a time: looking the image up would not give a bijection.""" - - groups = collections.defaultdict(collections.deque) - for j in sorted(columns): - form, sign = cls.jamp_column_form(columns[j]) - groups[form].append((j, sign)) - - action = {} - for j in sorted(columns): - image = dict((induced[i - 1] + 1, value) - for i, value in columns[j].items()) - form, sign = cls.jamp_column_form(image) - group = groups.get(form) - if not group: - return None - target, target_sign = group.popleft() - factor = sign * target_sign - other = columns[target] - if len(other) != len(image) or \ - any(other.get(i) != factor * value - for i, value in image.items()): - return None - action[j] = (target, factor) - return action - - def get_jamp_symmetry(self, matrix_element, all_element): - """Permutations leaving the JAMP matrix invariant: for each of them the - permutation of the color basis lines, and the permutation of the - amplitude columns with the sign that goes with it. None when there is - none, or when the matrix element does not carry a color basis.""" - - if not isinstance(matrix_element, helas_objects.HelasMatrixElement): - return None - color_basis = matrix_element.get('color_basis') - if not color_basis or len(color_basis) < 2: - return None - symmetry = color_amp.ColorBasisSymmetry(sorted(color_basis.keys())) - if not symmetry.generators1: - return None - - columns = collections.defaultdict(dict) - for (i, j), value in all_element.items(): - if value: - columns[j][i] = value - if not columns: - return None - - nb_line = len(symmetry.keys1) - rowperms, actions = [], [] - for induced in symmetry.generators1: - action = self.jamp_amp_permutation(columns, induced) - if action is None: - continue - rowperms.append([0] + [induced[i] + 1 for i in range(nb_line)]) - actions.append(action) - if not actions: - return None - - # one line per orbit is enough to see every sub-expression: any other - # line is the image of one of them, and so are the sub-expressions it - # holds. This is what keeps the scan below from being quadratic in the - # number of terms of the whole matrix. - parent = list(range(nb_line + 1)) - - def find(x): - while parent[x] != x: - parent[x] = parent[parent[x]] - x = parent[x] - return x - - for rowperm in rowperms: - for i in range(1, nb_line + 1): - ri, rj = find(i), find(rowperm[i]) - if ri != rj: - parent[ri] = rj - line_reps = [i for i in range(1, nb_line + 1) if find(i) == i] - - return {'rowperms': rowperms, 'actions': actions, - 'nb_line': nb_line, 'line_reps': line_reps} - - @staticmethod - def jamp_operation_image(action, operation): - """Image of the sub-expression operation=(j1,j2,R) under one - permutation, and the factor relating the column the image defines to - the image of the column operation defines.""" - - j1, j2, ratio = operation - first, sign1 = action[j1] - second, sign2 = action[j2] - if first < second: - return (first, second, ratio * sign2 / sign1), sign1 - return (second, first, sign1 / (sign2 * ratio)), sign2 * ratio - - def optimise_jamp_equivariant(self, all_element, symmetry): - """Same optimisation as optimise_jamp, but introducing whole orbits of - sub-expressions at a time so that the result is closed under the - symmetry. Fills self.jamp_orbits with, for every definition, the orbit - it belongs to and the definition and permutation it comes from.""" - - actions = [dict(action) for action in symmetry['actions']] - line_reps = symmetry['line_reps'] - added = 0 - defs = [] - # (orbit, parent definition, permutation) for every definition - tree = [] - # the definitions introduced together: none of them uses another, so - # they can be reordered freely - levels = [] - nb_orbit = 0 - - while True: - columns = collections.defaultdict(list) - lines = collections.defaultdict(list) - for (i, j), value in all_element.items(): - if value: - columns[j].append(i) - lines[i].append(j) - for line in lines.values(): - line.sort() - - # every sub-expression is the image of one living on a - # representative line, so only those have to be looked at - candidates = set() - for i in line_reps: - line = lines.get(i, []) - for pos, j1 in enumerate(line): - value = all_element[(i, j1)] - for j2 in line[pos + 1:]: - candidates.add((j1, j2, all_element[(i, j2)] / value)) - - max_count = 0 - best = [] - for operation in candidates: - count = len(self.jamp_operation_lines(all_element, columns, - operation)) - if count > max_count: - max_count, best = count, [operation] - elif count == max_count: - best.append(operation) - if max_count <= 1: - break - - orbits = self.jamp_operation_orbits(actions, best) - first_of_level = added + 1 - for orbit, parent in orbits: - rows = dict((operation, - self.jamp_operation_lines(all_element, columns, - operation)) - for operation in orbit) - if not self.jamp_orbit_usable(rows): - continue - index = {} - for operation in orbit: - added += 1 - index[operation] = added - origin, permutation = parent[operation] - tree.append((nb_orbit, index[origin] if origin else 0, - permutation)) - defs.append((added, operation[0], operation[1], - operation[2], len(rows[operation]))) - nb_orbit += 1 - for operation, new in index.items(): - j1, j2 = operation[0], operation[1] - for i in rows[operation]: - all_element[(i, -new)] = all_element[(i, j1)] - del all_element[(i, j1)] - del all_element[(i, j2)] - for action in actions: - for operation, new in index.items(): - image, factor = self.jamp_operation_image(action, - operation) - action[-new] = (-index[image], factor) - if added < first_of_level: - # nothing could be introduced as a whole orbit - break - levels.append((first_of_level, added)) - logger.log(5, "Define %d new shortcut reused %d times", - added - first_of_level + 1, max_count) - - self.jamp_orbits = {'tree': tree, 'nb_orbit': nb_orbit, - 'levels': levels, 'actions': actions, - 'symmetry': symmetry} - - if self.jamp_emit == 'tables' and self.jamp_greedy_tail: - # The orbit rounds stop while the JAMP lines still hold a good many - # terms, since an orbit can only be taken as a whole. The plain - # scan has no such scruple and can still shorten those lines. Its - # sub-expressions are not orbits of anything, which rules them out - # of the recipes, but the table emission does not care: there a - # definition costs three numbers of DATA and one indirect add, - # against a term of a line and a direct add. - all_element, tail = self.optimise_jamp(all_element, added=added) - defs.extend(tail) - - return all_element, defs - - @staticmethod - def jamp_operation_lines(all_element, columns, operation): - """Lines where both columns of the sub-expression are still there with - its ratio. The values are read from the matrix as it is now, so lines - already taken by an orbit introduced before are simply gone.""" - - j1, j2, ratio = operation - res = [] - for i in columns.get(j1, ()): - value = all_element.get((i, j1), 0) - if not value: - continue - other = all_element.get((i, j2), 0) - if other and other / value == ratio: - res.append(i) - return res - - def jamp_operation_orbits(self, actions, operations): - """Orbits of the sub-expressions, walked breadth first, with the - (sub-expression, permutation) each of them is reached from.""" - - seen = set() - orbits = [] - for start in sorted(operations, key=lambda op: (op[0], op[1], - op[2].real, - op[2].imag)): - if start in seen: - continue - orbit, parent = [start], {start: (None, 0)} - seen.add(start) - queue = collections.deque([start]) - while queue: - current = queue.popleft() - for position, action in enumerate(actions): - image = self.jamp_operation_image(action, current)[0] - if image in seen: - continue - seen.add(image) - parent[image] = (current, position + 1) - orbit.append(image) - queue.append(image) - orbits.append((orbit, parent)) - return orbits - - @staticmethod - def jamp_i_power(factor): - """The exponent of i this factor is, or None when it is not one of the - four powers of i. The factors the optimisation produces are products of - signs and of the i the color coefficients carry, so this is what they - all are in practice.""" - - value = complex(factor) - for exponent, power in enumerate((1, 1j, -1, -1j)): - if value == power: - return exponent - return None - def jamp_orbit_recipes(self, defs, nb_amp): """Describe the definitions by one recipe per orbit: the amplitude permutations, the first definition of every orbit, and the definitions @@ -3804,30 +3487,6 @@ def jamp_orbit_tables(self, defs, nb_amp): 'complex_factor': any(complex(new_defs[one - 1][3]).imag for one in general)} - @staticmethod - def jamp_definition_levels(defs): - """Group the definitions by how deep they sit in their own operands: - one which uses no other is at the first level, and any other one comes - after both of the ones it uses. Nothing inside a level uses anything - else of that level, so they can be reordered freely. - - Read off the operands rather than off the rounds of the optimisation, - so that whatever the plain scan adds at the end lands where it belongs. - The operands of a definition always come before it, so one pass is - enough.""" - - depth = {} - levels = collections.defaultdict(list) - for index, left, right, _ratio, _count in defs: - here = 0 - if left < 0: - here = max(here, depth[-left]) - if right < 0: - here = max(here, depth[-right]) - depth[index] = here + 1 - levels[here + 1].append(index) - return [levels[key] for key in sorted(levels)] - @staticmethod def jamp_number_data_lines(name, values, per_line, var='IJMP'): """DATA statements filling one array with the given constants.""" @@ -4117,10 +3776,18 @@ def jamp_tables_allowed(self): return True + def jamp_greedy_tail_enabled(self): + """The tail of the plain scan is only within reach of the emission + which writes the operands out: the definitions it adds are not orbits + of anything, so INIT_JAMP could not rebuild them from the recipes.""" + + return self.jamp_greedy_tail and self.jamp_emit == 'tables' + def jamp_orbit_allowed(self, matrix_element): - """Whether the orbit equivariant optimisation is used here.""" + """Whether the orbit equivariant optimisation is used here: only the + templates which know how to write the definitions it produces.""" - if not self.jamp_orbit: + if not super().jamp_orbit_allowed(matrix_element): return False if isinstance(self, ProcessExporterFortranME): @@ -4269,31 +3936,6 @@ def jamp_orbit_reach(actions, chosen, first): queue.append(image) return len(seen) - @staticmethod - def jamp_orbit_usable(rows): - """Restrict an orbit to the entries only one of its sub-expressions - wants, and say whether what is left can be introduced as a whole. Which - of two sub-expressions of the same orbit gets a shared entry cannot be - decided in a way that commutes with the symmetry, so those entries are - left in the matrix and get another chance in a later round.""" - - sizes = set(len(use) for use in rows.values()) - if len(sizes) != 1 or sizes == set([0]): - return False - entry = collections.Counter() - for operation, use in rows.items(): - for i in use: - entry[(i, operation[0])] += 1 - entry[(i, operation[1])] += 1 - if max(entry.values()) == 1: - return True - for operation in list(rows): - rows[operation] = [i for i in rows[operation] - if entry[(i, operation[0])] == 1 - and entry[(i, operation[1])] == 1] - sizes = set(len(use) for use in rows.values()) - return len(sizes) == 1 and sizes != set([0]) - def get_pdf_lines(self, matrix_element, ninitial, subproc_group = False, vector=False): diff --git a/madgraph/iolibs/jamp_optimiser.py b/madgraph/iolibs/jamp_optimiser.py index f11a7a0208..2b985c9bd1 100644 --- a/madgraph/iolibs/jamp_optimiser.py +++ b/madgraph/iolibs/jamp_optimiser.py @@ -24,10 +24,18 @@ repeated pieces by shared sub-expressions, so that the matrix is left with far fewer entries and a list of definitions to compute first. +Two searches are here, and optimise_jamp_best picks between them: the plain +greedy scan, which takes whatever sub-expression is worth most at each step, +and the orbit equivariant one, which only takes whole orbits of the +permutations leaving the color basis invariant and so leaves the matrix +invariant at every step. + Nothing here knows about fortran or C++: it takes the coefficient matrix and gives back the reduced matrix and the definitions. The exporters print that in -their own language (get_JAMP_lines for fortran, get_jamp_accumulation_lines for -the C++/cudacpp writer). +their own language (get_JAMP_lines for fortran, build_jamp_plan for the +C++/cudacpp writer), and how they print it is what decides which of the +optimisations they may use -- see jamp_orbit_allowed and +jamp_greedy_tail_enabled. """ from __future__ import absolute_import @@ -38,6 +46,8 @@ import logging import time +import madgraph.core.color_amp as color_amp +import madgraph.core.helas_objects as helas_objects import madgraph.various.banner as banner_mod logger = logging.getLogger('madgraph.export_v4') @@ -46,9 +56,7 @@ class JampOptimiser(object): """The common sub-expression search over the JAMP coefficient matrix. - Mixed into the exporters, which supply the printing. A subclass that hands - a symmetry to optimise_jamp must also provide optimise_jamp_best (only the - fortran exporter does, see export_v4).""" + Mixed into the exporters, which supply the printing.""" # Off by default: the plain output of a backend is the expanded one, and # each exporter switches this on for itself. 'jamp_optim' in cmd_options @@ -59,6 +67,20 @@ class JampOptimiser(object): # take the power of i shared by every coefficient out before searching, so # that the search walks over whole numbers (see optimise_jamp_matrix) jamp_integer_walk = True + # Introduce the sub-expressions by whole orbits of the permutations leaving + # the color basis invariant instead of one at a time (see + # optimise_jamp_equivariant). Off by default, each backend switches it on + # for itself where it measured a gain. + jamp_orbit = False + # finish with the plain scan once the orbit rounds have nothing left to + # take as a whole (see jamp_greedy_tail_enabled) + jamp_greedy_tail = True + # up to this many entries in the matrix, both optimisations are run and the + # shorter result kept (see optimise_jamp_best) + jamp_compare_max_size = 20000 + # what the orbit rounds did, for an emission which wants to describe the + # definitions by one recipe per orbit rather than one by one + jamp_orbits = None def jamp_optim_enabled(self): """Whether to run the optimisation, --jamp_optim first.""" @@ -138,7 +160,8 @@ def jamp_apply_phase(new_mat, phase): for key in new_mat: new_mat[key] = new_mat[key] * phase - def optimise_jamp_matrix(self, all_element, symmetry=None): + def optimise_jamp_matrix(self, all_element, symmetry=None, + matrix_element=None): """Run the optimisation over the coefficient matrix and return (new_mat, defs): - defs is a list of (i, op1, op2, frac, nb): definition number i is @@ -148,9 +171,13 @@ def optimise_jamp_matrix(self, all_element, symmetry=None): input except that a negative amplitude index means the definition of that number. + With a matrix element and jamp_orbit_allowed saying so, the color basis + symmetry is read off the matrix and the sub-expressions are introduced + by whole orbits of it; the orbits are left in self.jamp_orbits. + all_element is consumed (the optimisation works in place). The fortran - exporter runs the three steps itself, since it has to look at the - walked matrix to work out the color basis symmetry in between.""" + exporter runs the three steps itself, since it has its own way of + finding the matrix element the color amplitudes came from.""" if len(all_element) > 1000: logger.info("Computing Color-Flow optimization [%s term]", @@ -160,7 +187,13 @@ def optimise_jamp_matrix(self, all_element, symmetry=None): start_time = 0 self.myjamp_count = 0 + self.jamp_orbits = None phase = self.jamp_walk_integers(all_element) + # the symmetry is read off the matrix once the phase is out of it, so + # that the columns compare as whole numbers + if symmetry is None and matrix_element is not None and \ + self.jamp_orbit_allowed(matrix_element): + symmetry = self.get_jamp_symmetry(matrix_element, all_element) new_mat, defs = self.optimise_jamp(all_element, symmetry=symmetry) self.jamp_apply_phase(new_mat, phase) if start_time: @@ -379,3 +412,382 @@ def jamp_definition_order(defs): rank[i] = position order = sorted(ready, key=lambda i: (ready[i], rank[i])) return order, ready + + def optimise_jamp_best(self, all_element, symmetry): + """Taking whole orbits only pays once there is enough of them to share: + on a small matrix it can end up asking for more additions than the plain + scan, which is free to take whatever it likes. g g > t t~ g is such a + case, 46 additions against 39. + + Small matrices are cheap to optimise, so rather than guess where the + turn is, do both and keep the shorter. Above that size only the orbit + version is run: it wins by a wide margin on everything that big, and + the plain scan is the slow one there.""" + + orbit_element, orbit_defs = self.optimise_jamp_equivariant( + dict(all_element), symmetry) + if len(all_element) > self.jamp_compare_max_size: + return orbit_element, orbit_defs + + orbits = self.jamp_orbits + plain_element, plain_defs = self.optimise_jamp(dict(all_element)) + if self.jamp_operation_count(plain_element, plain_defs) < \ + self.jamp_operation_count(orbit_element, orbit_defs): + self.jamp_orbits = None + return plain_element, plain_defs + self.jamp_orbits = orbits + return orbit_element, orbit_defs + + def jamp_orbit_allowed(self, matrix_element): + """Whether the orbit equivariant optimisation is used here, + --jamp_orbit first. A backend which only accepts it for some of the + templates it writes says so by overriding this.""" + + cmd_options = getattr(self, 'cmd_options', None) or {} + if 'jamp_orbit' in cmd_options: + return banner_mod.ConfigFile.format_variable( + cmd_options['jamp_orbit'], bool, 'jamp_orbit') + return self.jamp_orbit + + def jamp_greedy_tail_enabled(self): + """Whether the orbit rounds are finished off by the plain scan. What + the tail adds are ordinary sub-expressions, not orbits of anything, so + an emission which rebuilds the definitions from one recipe per orbit + cannot describe them; one which writes them down can.""" + + return self.jamp_greedy_tail + + #=========================================================================== + # Orbit equivariant version of the JAMP optimisation + #=========================================================================== + # A permutation of the external color indices which maps the color basis + # onto itself (see color_amp.ColorBasisSymmetry) also permutes the columns + # of the JAMP matrix, up to a sign. The whole matrix is then invariant, so + # the sub-expressions the optimisation looks for come in orbits: every one + # of them is worth exactly as much as the others. Introducing a whole orbit + # at a time, rather than one sub-expression at a time as the plain scan + # does, leaves the matrix invariant at every step, and the definitions can + # be written as one recipe per orbit. + + @staticmethod + def jamp_column_form(column): + """Canonical form of one column of the JAMP matrix up to a global sign, + together with the sign which was taken out.""" + + entries = sorted(column.items()) + first = entries[0][1] + sign = -1 if (first.real, first.imag) < (0., 0.) else 1 + return tuple((i, sign * value) for i, value in entries), sign + + @classmethod + def jamp_amp_permutation(cls, columns, induced): + """Permutation of the amplitudes induced by the permutation induced of + the color basis: return {amp: (amp, sign)} such that + + M[induced[i], sigma(j)] = sign(j) * M[i, j] + + or None if the columns are not mapped onto each other. + + Several amplitudes often have the very same column, so the columns are + gathered by their canonical form and one target is taken out of each + group at a time: looking the image up would not give a bijection.""" + + groups = collections.defaultdict(collections.deque) + for j in sorted(columns): + form, sign = cls.jamp_column_form(columns[j]) + groups[form].append((j, sign)) + + action = {} + for j in sorted(columns): + image = dict((induced[i - 1] + 1, value) + for i, value in columns[j].items()) + form, sign = cls.jamp_column_form(image) + group = groups.get(form) + if not group: + return None + target, target_sign = group.popleft() + factor = sign * target_sign + other = columns[target] + if len(other) != len(image) or \ + any(other.get(i) != factor * value + for i, value in image.items()): + return None + action[j] = (target, factor) + return action + + def get_jamp_symmetry(self, matrix_element, all_element): + """Permutations leaving the JAMP matrix invariant: for each of them the + permutation of the color basis lines, and the permutation of the + amplitude columns with the sign that goes with it. None when there is + none, or when the matrix element does not carry a color basis.""" + + if not isinstance(matrix_element, helas_objects.HelasMatrixElement): + return None + color_basis = matrix_element.get('color_basis') + if not color_basis or len(color_basis) < 2: + return None + symmetry = color_amp.ColorBasisSymmetry(sorted(color_basis.keys())) + if not symmetry.generators1: + return None + + columns = collections.defaultdict(dict) + for (i, j), value in all_element.items(): + if value: + columns[j][i] = value + if not columns: + return None + + nb_line = len(symmetry.keys1) + rowperms, actions = [], [] + for induced in symmetry.generators1: + action = self.jamp_amp_permutation(columns, induced) + if action is None: + continue + rowperms.append([0] + [induced[i] + 1 for i in range(nb_line)]) + actions.append(action) + if not actions: + return None + + # one line per orbit is enough to see every sub-expression: any other + # line is the image of one of them, and so are the sub-expressions it + # holds. This is what keeps the scan below from being quadratic in the + # number of terms of the whole matrix. + parent = list(range(nb_line + 1)) + + def find(x): + while parent[x] != x: + parent[x] = parent[parent[x]] + x = parent[x] + return x + + for rowperm in rowperms: + for i in range(1, nb_line + 1): + ri, rj = find(i), find(rowperm[i]) + if ri != rj: + parent[ri] = rj + line_reps = [i for i in range(1, nb_line + 1) if find(i) == i] + + return {'rowperms': rowperms, 'actions': actions, + 'nb_line': nb_line, 'line_reps': line_reps} + + @staticmethod + def jamp_operation_image(action, operation): + """Image of the sub-expression operation=(j1,j2,R) under one + permutation, and the factor relating the column the image defines to + the image of the column operation defines.""" + + j1, j2, ratio = operation + first, sign1 = action[j1] + second, sign2 = action[j2] + if first < second: + return (first, second, ratio * sign2 / sign1), sign1 + return (second, first, sign1 / (sign2 * ratio)), sign2 * ratio + + def optimise_jamp_equivariant(self, all_element, symmetry): + """Same optimisation as optimise_jamp, but introducing whole orbits of + sub-expressions at a time so that the result is closed under the + symmetry. Fills self.jamp_orbits with, for every definition, the orbit + it belongs to and the definition and permutation it comes from.""" + + actions = [dict(action) for action in symmetry['actions']] + line_reps = symmetry['line_reps'] + added = 0 + defs = [] + # (orbit, parent definition, permutation) for every definition + tree = [] + # the definitions introduced together: none of them uses another, so + # they can be reordered freely + levels = [] + nb_orbit = 0 + + while True: + columns = collections.defaultdict(list) + lines = collections.defaultdict(list) + for (i, j), value in all_element.items(): + if value: + columns[j].append(i) + lines[i].append(j) + for line in lines.values(): + line.sort() + + # every sub-expression is the image of one living on a + # representative line, so only those have to be looked at + candidates = set() + for i in line_reps: + line = lines.get(i, []) + for pos, j1 in enumerate(line): + value = all_element[(i, j1)] + for j2 in line[pos + 1:]: + candidates.add((j1, j2, all_element[(i, j2)] / value)) + + max_count = 0 + best = [] + for operation in candidates: + count = len(self.jamp_operation_lines(all_element, columns, + operation)) + if count > max_count: + max_count, best = count, [operation] + elif count == max_count: + best.append(operation) + if max_count <= 1: + break + + orbits = self.jamp_operation_orbits(actions, best) + first_of_level = added + 1 + for orbit, parent in orbits: + rows = dict((operation, + self.jamp_operation_lines(all_element, columns, + operation)) + for operation in orbit) + if not self.jamp_orbit_usable(rows): + continue + index = {} + for operation in orbit: + added += 1 + index[operation] = added + origin, permutation = parent[operation] + tree.append((nb_orbit, index[origin] if origin else 0, + permutation)) + defs.append((added, operation[0], operation[1], + operation[2], len(rows[operation]))) + nb_orbit += 1 + for operation, new in index.items(): + j1, j2 = operation[0], operation[1] + for i in rows[operation]: + all_element[(i, -new)] = all_element[(i, j1)] + del all_element[(i, j1)] + del all_element[(i, j2)] + for action in actions: + for operation, new in index.items(): + image, factor = self.jamp_operation_image(action, + operation) + action[-new] = (-index[image], factor) + if added < first_of_level: + # nothing could be introduced as a whole orbit + break + levels.append((first_of_level, added)) + logger.log(5, "Define %d new shortcut reused %d times", + added - first_of_level + 1, max_count) + + self.jamp_orbits = {'tree': tree, 'nb_orbit': nb_orbit, + 'levels': levels, 'actions': actions, + 'symmetry': symmetry} + + if self.jamp_greedy_tail_enabled(): + # The orbit rounds stop while the JAMP lines still hold a good many + # terms, since an orbit can only be taken as a whole. The plain + # scan has no such scruple and can still shorten those lines. Its + # sub-expressions are not orbits of anything, so an emission which + # rebuilds them from the recipes cannot have them (see + # jamp_greedy_tail_enabled), but one writing them down does not + # care. + all_element, tail = self.optimise_jamp(all_element, added=added) + defs.extend(tail) + + return all_element, defs + + @staticmethod + def jamp_operation_lines(all_element, columns, operation): + """Lines where both columns of the sub-expression are still there with + its ratio. The values are read from the matrix as it is now, so lines + already taken by an orbit introduced before are simply gone.""" + + j1, j2, ratio = operation + res = [] + for i in columns.get(j1, ()): + value = all_element.get((i, j1), 0) + if not value: + continue + other = all_element.get((i, j2), 0) + if other and other / value == ratio: + res.append(i) + return res + + def jamp_operation_orbits(self, actions, operations): + """Orbits of the sub-expressions, walked breadth first, with the + (sub-expression, permutation) each of them is reached from.""" + + seen = set() + orbits = [] + for start in sorted(operations, key=lambda op: (op[0], op[1], + op[2].real, + op[2].imag)): + if start in seen: + continue + orbit, parent = [start], {start: (None, 0)} + seen.add(start) + queue = collections.deque([start]) + while queue: + current = queue.popleft() + for position, action in enumerate(actions): + image = self.jamp_operation_image(action, current)[0] + if image in seen: + continue + seen.add(image) + parent[image] = (current, position + 1) + orbit.append(image) + queue.append(image) + orbits.append((orbit, parent)) + return orbits + + @staticmethod + def jamp_i_power(factor): + """The exponent of i this factor is, or None when it is not one of the + four powers of i. The factors the optimisation produces are products of + signs and of the i the color coefficients carry, so this is what they + all are in practice.""" + + value = complex(factor) + for exponent, power in enumerate((1, 1j, -1, -1j)): + if value == power: + return exponent + return None + + @staticmethod + def jamp_definition_levels(defs): + """Group the definitions by how deep they sit in their own operands: + one which uses no other is at the first level, and any other one comes + after both of the ones it uses. Nothing inside a level uses anything + else of that level, so they can be reordered freely. + + Read off the operands rather than off the rounds of the optimisation, + so that whatever the plain scan adds at the end lands where it belongs. + The operands of a definition always come before it, so one pass is + enough.""" + + depth = {} + levels = collections.defaultdict(list) + for index, left, right, _ratio, _count in defs: + here = 0 + if left < 0: + here = max(here, depth[-left]) + if right < 0: + here = max(here, depth[-right]) + depth[index] = here + 1 + levels[here + 1].append(index) + return [levels[key] for key in sorted(levels)] + + @staticmethod + def jamp_orbit_usable(rows): + """Restrict an orbit to the entries only one of its sub-expressions + wants, and say whether what is left can be introduced as a whole. Which + of two sub-expressions of the same orbit gets a shared entry cannot be + decided in a way that commutes with the symmetry, so those entries are + left in the matrix and get another chance in a later round.""" + + sizes = set(len(use) for use in rows.values()) + if len(sizes) != 1 or sizes == set([0]): + return False + entry = collections.Counter() + for operation, use in rows.items(): + for i in use: + entry[(i, operation[0])] += 1 + entry[(i, operation[1])] += 1 + if max(entry.values()) == 1: + return True + for operation in list(rows): + rows[operation] = [i for i in rows[operation] + if entry[(i, operation[0])] == 1 + and entry[(i, operation[1])] == 1] + sizes = set(len(use) for use in rows.values()) + return len(sizes) == 1 and sizes != set([0]) diff --git a/madmatrix/model_handling.py b/madmatrix/model_handling.py index 3eae0a3e52..0177684518 100644 --- a/madmatrix/model_handling.py +++ b/madmatrix/model_handling.py @@ -2424,6 +2424,10 @@ class MadMatrixUFOHelasCallWriter(helas_call_writers.GPUFOHelasCallWriter, # optimisation finds, instead of one line per (color flow, amplitude) pair # (see build_jamp_plan). Toggled by --jamp_optim=True|False. jamp_optim = True + # Look for those sub-expressions by whole orbits of the permutations + # leaving the color basis invariant (see JampOptimiser). Toggled by + # --jamp_orbit=True|False. + jamp_orbit = True # Class structure information # - object # - dict(object) [built-in] @@ -2664,7 +2668,7 @@ 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 build_jamp_plan(self, color_amplitudes): + 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): - captures[n] are the lines to write while amplitude n sits in @@ -2681,7 +2685,8 @@ def build_jamp_plan(self, color_amplitudes): all_element = self.jamp_matrix(color_amplitudes) if not all_element: return None - new_mat, defs = self.optimise_jamp_matrix(all_element) + new_mat, defs = self.optimise_jamp_matrix(all_element, + matrix_element=matrix_element) if not defs: return None order, ready = self.jamp_definition_order(defs) @@ -2768,7 +2773,7 @@ def super_get_matrix_element_calls(self, matrix_element, color_amplitudes, multi color[namp][njamp] = coeff # Color flows through shared sub-expressions (None to write them out # one (color flow, amplitude) pair at a time, as before) - jamp_plan = self.build_jamp_plan(color_amplitudes) + jamp_plan = self.build_jamp_plan(matrix_element, color_amplitudes) self.nb_tmp_jamp = jamp_plan[0] if jamp_plan else 0 if jamp_plan is not None: _ntmp, jamp_captures, jamp_combines, jamp_final = jamp_plan From 467862fe5bf94e874554411a5855342497be475b Mon Sep 17 00:00:00 2001 From: Olivier Mattelaer Date: Sun, 9 Aug 2026 21:28:05 +0200 Subject: [PATCH 41/42] mg7: pick the color flow from a mask built on the color flow basis The mg7 exporter can run its color sum on the (n-2)! Del Duca-Dixon-Maltoni structures while a color flow is still picked among the (n-1)! trace ones, but set_channels_colors_map built active_colors -- the mask which becomes icolamp -- by walking self.color_basis, i.e. the basis of the color sum. The C++ selection walks that mask over ncolor_flow entries, so a DDM build declared icolamp[nconfig][ncolor] and read it as icolamp[nconfig][ncolor_flow]: for g g > g g a [3][2] array read 6 wide, which reinterprets the neighbouring rows and then runs off the end of the object. Reading icolamp the way the selection does gave {1,2,4,5}, {0,2,3}, {0,1,5} for the three configs instead of {0,1,3,5}, {1,2,3,4}, {0,2,4,5}. |M|^2 does not depend on any of this, which is why the DDM port passed every check run so far. Events do: over 3 x 100k g g > g g events the color flow fractions came out (9.35, 9.38, 31.38, 9.15, 31.30, 9.44)% on the trace basis against (10.54, 1.39, 32.32, 12.83, 32.46, 10.46)% on the DDM one -- 134 sigma on the second flow -- at an unchanged cross section (1.2 sigma), since a wrong color selection never moves the weight. Give the exporter a color_flow_basis next to its color_basis and build active_colors, the color_flows table and both nb_color values from it. The two bases now write the same icolamp, and the same processes agree flow by flow over 3 x 100k events each: g g > g g worst 2.0 sigma over 6 flows (xsec 0.1), g g > g g g worst 1.9 sigma over 24 (xsec 0.5), and u u~ > g g, which never leaves the trace basis, generates identical source and agrees at 0.1 sigma. The GPU select_col_and_diag kernel had the same ncolor/ncolor_flow mismatch and is fixed the same way; there is no CUDA toolchain here, so that path is generated but not compiled. Co-Authored-By: Claude Opus 5 --- madgraph/iolibs/export_mg7.py | 16 +++++- .../process_function_definitions.inc | 10 ++-- madmatrix/model_handling.py | 10 +++- tests/unit_tests/iolibs/test_export_cpp.py | 55 +++++++++++++++++++ 4 files changed, 81 insertions(+), 10 deletions(-) diff --git a/madgraph/iolibs/export_mg7.py b/madgraph/iolibs/export_mg7.py index a9b5e458c3..5985420760 100644 --- a/madgraph/iolibs/export_mg7.py +++ b/madgraph/iolibs/export_mg7.py @@ -25,6 +25,13 @@ def __init__(self, matrix_element, cpp_helas_call_writer): self.process = self.amplitude.get("process") self.legs = self.process.get("legs_with_decays") self.color_basis = self.matrix_element.get("color_basis") + # The basis a color flow is picked among: always the trace one, which + # is the color basis itself unless the color sum runs on the DDM basis. + # Everything indexing a color flow -- the color_flows table, the + # active_colors masks, icolamp -- has to use this one and not the + # (smaller) basis of the color sum. + self.color_flow_basis = self.color_basis.get_flow_basis() \ + if self.color_basis else self.color_basis self.set_topology() self.set_flavor_indices() self.set_active_flavors() @@ -104,9 +111,12 @@ def set_active_flavors(self): def set_channels_colors_map(self): if self.color_basis: + # active_colors ends up in the icolamp mask, which is walked over + # the color flows, so it must be indexed on the flow basis + flow_basis = self.color_flow_basis diag_jamps = defaultdict(list) - for ijamp, col_basis_elem in enumerate(sorted(self.color_basis.keys())): - for diag_tuple in self.color_basis[col_basis_elem]: + for ijamp, col_basis_elem in enumerate(sorted(flow_basis.keys())): + for diag_tuple in flow_basis[col_basis_elem]: diag_jamps[diag_tuple[0]].append(ijamp) sym_indices, sym_perms, _ = find_symmetry( @@ -208,7 +218,7 @@ def get_subprocess_info(self, proc_dir, lib_me_path): # Get the list of color flows. This is about color flows, so # always the trace basis, even when the color sum runs on the DDM # one. - color_flow_dicts = self.color_basis.get_flow_basis().\ + color_flow_dicts = self.color_flow_basis.\ color_flow_decomposition(repr_dict, n_initial) # And output them properly color_flows = [ diff --git a/madgraph/iolibs/template_files/madmatrix/process_function_definitions.inc b/madgraph/iolibs/template_files/madmatrix/process_function_definitions.inc index eaef9990c8..1ae25f9b05 100644 --- a/madgraph/iolibs/template_files/madmatrix/process_function_definitions.inc +++ b/madgraph/iolibs/template_files/madmatrix/process_function_definitions.inc @@ -815,7 +815,7 @@ namespace mg5amcCpu fptype_sv jamp2_sv[ncolor_flow] = { 0 }; assert( allJamp2s != nullptr ); // sanity check using J2_ACCESS = DeviceAccessJamp2; - for( int icolC = 0; icolC < ncolor; icolC++ ) + for( int icolC = 0; icolC < ncolor_flow; icolC++ ) jamp2_sv[icolC] = J2_ACCESS::kernelAccessIcolConst( allJamp2s, icolC ); // NB (see #877): in the array channel2iconfig, the input index uses C indexing (channelId -1), the output index uses F indexing (iconfig) // NB (see #917): mgOnGpu::channel2iconfig returns an int (which may be -1), not an unsigned int! @@ -830,9 +830,9 @@ namespace mg5amcCpu printf( "INTERNAL ERROR! Cannot choose an event-by-event random color for channelId=%%d (invalid SDE iconfig=%%d\n > nconfig=%%d)", channelId, iconfig, mgOnGpu::nconfigSDE ); assert( iconfig <= (int)mgOnGpu::nconfigSDE ); // SANITY CHECK #917 } - fptype targetamp[ncolor] = { 0 }; + fptype targetamp[ncolor_flow] = { 0 }; // NB (see #877): explicitly use 'icolC' rather than 'icol' to indicate that icolC uses C indexing in [0, N_colors-1] - for( int icolC = 0; icolC < ncolor; icolC++ ) + for( int icolC = 0; icolC < ncolor_flow; icolC++ ) { if( icolC == 0 ) targetamp[icolC] = 0; @@ -842,9 +842,9 @@ namespace mg5amcCpu if( mgOnGpu::icolamp[iconfig - 1][icolC] ) targetamp[icolC] += jamp2_sv[icolC]; } //printf( "sigmaKin: ievt=%%4d rndcol=%%f\n", ievt, allrndcol[ievt] ); - for( int icolC = 0; icolC < ncolor; icolC++ ) + for( int icolC = 0; icolC < ncolor_flow; icolC++ ) { - if( allrndcol[ievt] < ( targetamp[icolC] / targetamp[ncolor - 1] ) ) + if( allrndcol[ievt] < ( targetamp[icolC] / targetamp[ncolor_flow - 1] ) ) { allselcol[ievt] = icolC + 1; // NB Fortran [1,ncolor], cudacpp [0,ncolor-1] //printf( "sigmaKin: ievt=%%d icol=%%d\n", ievt, icolC+1 ); diff --git a/madmatrix/model_handling.py b/madmatrix/model_handling.py index 0177684518..320c8a53c4 100644 --- a/madmatrix/model_handling.py +++ b/madmatrix/model_handling.py @@ -1783,7 +1783,9 @@ def get_sigmaKin_lines(self, color_amplitudes, write=True): replace_dict['madE_update_answer'] = ' allMEs[iproc*nprocesses + ievt] *= multi_chanel_num/multi_chanel_denom;' replace_dict['nb_channel'] = len(self.multi_channel_map) - replace_dict['nb_color'] = max(1, len(self.matrix_elements[0].get('color_basis'))) + # same meaning as in edit_coloramps: the number of color flows, which + # is not the size of the color basis when the color sum runs on the DDM one + replace_dict['nb_color'] = max(1, len(self.color_flow_basis)) replace_dict['cpp_blas_helicity_loop'] = '' replace_dict['cpp_blas_helicity_loop_end'] = '' @@ -2111,7 +2113,11 @@ def edit_coloramps(self): replace_dict['nb_channel'] = len(self.active_color_map) # here I can do the conversion in between the active color map and true, false, and obtain a C++ compatible thing immediately replace_dict['nb_diag'] = nb_diag - nb_color = max(1, len(self.color_basis)) + # icolamp is the mask of the color flows allowed for a config, and the + # selection walks it over ncolor_flow entries, so it is dimensioned on + # the flow basis -- which is larger than the color basis itself when + # the color sum runs on the DDM one + nb_color = max(1, len(self.color_flow_basis)) replace_dict['nb_color'] = nb_color # AV extra formatting (e.g. gg_tt was "{{true,true};,{true,false};,{false,true};};") ###misc.sprint(replace_dict['is_LC']) diff --git a/tests/unit_tests/iolibs/test_export_cpp.py b/tests/unit_tests/iolibs/test_export_cpp.py index d621b4dd59..4c24d68fb8 100755 --- a/tests/unit_tests/iolibs/test_export_cpp.py +++ b/tests/unit_tests/iolibs/test_export_cpp.py @@ -28,6 +28,7 @@ import aloha.create_aloha as create_aloha import madgraph.iolibs.export_cpp as export_cpp +import madgraph.iolibs.export_mg7 as export_mg7 import madgraph.iolibs.export_v4 as export_v4 import madgraph.iolibs.file_writers as writers import madgraph.iolibs.helas_call_writers as helas_call_writer @@ -37,6 +38,7 @@ import madgraph.core.base_objects as base_objects import madgraph.core.color_algebra as color +import madgraph.core.color_amp as color_amp import madgraph.core.helas_objects as helas_objects import madgraph.core.diagram_generation as diagram_generation @@ -929,3 +931,56 @@ def test_cpp_export_decay_chain_broken_symmetry_metadata(self): self.assertIn('const int n_components = 3;', rendered) self.assertIn('const int comp_old[n_components] = {1,1,1};', rendered) self.assertIn('const int block_len[n_entries] = {2,2,1,1,1,1};', rendered) + + +#=============================================================================== +# DDMColorFlowMG7Test +#=============================================================================== +class DDMColorFlowMG7Test(unittest.TestCase): + """The mg7 exporter picks a color flow among the (n-1)! trace structures + even when the color sum runs on the (n-2)! DDM ones, so everything which + indexes a color flow must be built on the trace basis. Switching the color + basis changes how the jamps are computed, never which color flows exist, + so none of it may depend on the mode.""" + + def get_exporter(self, ids, ddm): + """The mg7 exporter for the all-gluon process with npar = len(ids), + built with or without the DDM color basis.""" + + color_amp.set_ddm_basis(ddm, with_flow=ddm) + try: + model = import_ufo.import_model('sm') + legs = base_objects.LegList( + [base_objects.Leg({'id': pdg, 'state': i > 1}) + for i, pdg in enumerate(ids)]) + amplitude = diagram_generation.Amplitude( + base_objects.Process({'legs': legs, 'model': model})) + matrix_element = helas_objects.HelasMatrixElement(amplitude) + return export_mg7.OneProcessExporterMG7( + matrix_element, helas_call_writer.CPPUFOHelasCallWriter(model)) + finally: + color_amp.set_ddm_basis(False) + + def test_ddm_color_flow_basis_is_the_trace_one(self): + """The color sum shrinks to (n-2)! structures, the color flows stay + the (n-1)! trace ones.""" + + for npar, ncolor, nflow in [(4, 2, 6), (5, 6, 24)]: + exporter = self.get_exporter([21] * npar, ddm=True) + self.assertEqual(len(exporter.color_basis), ncolor) + self.assertEqual(len(exporter.color_flow_basis), nflow) + + def test_ddm_active_colors_index_the_color_flows(self): + """active_colors ends up in the icolamp mask, which the color + selection walks over ncolor_flow entries: it must be the same mask + the trace basis writes, not one indexed on the smaller DDM basis.""" + + for npar in (4, 5): + trace = self.get_exporter([21] * npar, ddm=False) + ddm = self.get_exporter([21] * npar, ddm=True) + self.assertEqual(ddm.active_color_map, trace.active_color_map) + # and it is a mask over the color flows, not over the color sum + nflow = len(ddm.color_flow_basis) + for active_colors in ddm.active_color_map: + self.assertTrue(active_colors) + self.assertLess(max(active_colors), nflow) From 9f984d9bf62e8be8b9f01e556332df58b4253cbb Mon Sep 17 00:00:00 2001 From: Olivier Mattelaer Date: Sun, 9 Aug 2026 23:26:08 +0200 Subject: [PATCH 42/42] hel_recycle: never wrap a statement onto a blank physical line do_multiline breaks a long statement at the last space which fits in 72 characters. When the only such space is the statement's own indentation the first chunk is all blanks, so the emitted physical line is empty and the continuation which follows it is attached by the compiler to the *previous* statement. g g > g g g on the DDM color basis is the first process to hit it: the Kleiss-Kuijf lines rebuilding the trace JAMPs are long and contain no space of their own, and matrix1_optim.f came out as JAMPF(1,1)=+2D0*(+IMAG1*JAMP(6,1)) $ JAMPF(2,1)=+2D0*(-IMAG1*JAMP(3,1)-...) which gfortran rejects with "Unclassifiable statement", so the madevent run dies in the compilation step. Break mid-token in that case, which is what the no-space-at-all branch next to it already does. The condition only fires where the old code emitted a blank chunk, so any matrix element which built before is byte for byte what it was; of the arms generated here only g g > g g g on the DDM basis contained one. With this, g g > g g g builds on both color bases and the two agree exactly: same cross section, and all 10000 color columns of all 2000 events identical at equal seed. Co-Authored-By: Claude Opus 5 --- madgraph/madevent/hel_recycle.py | 7 +- tests/unit_tests/madevent/test_hel_recycle.py | 81 +++++++++++++++++++ 2 files changed, 87 insertions(+), 1 deletion(-) create mode 100644 tests/unit_tests/madevent/test_hel_recycle.py diff --git a/madgraph/madevent/hel_recycle.py b/madgraph/madevent/hel_recycle.py index af0a1e7daa..60cf2b20f7 100755 --- a/madgraph/madevent/hel_recycle.py +++ b/madgraph/madevent/hel_recycle.py @@ -960,7 +960,12 @@ def do_multiline(line): remaining = line while len(remaining) > char_limit: split_at = remaining.rfind(' ', 0, char_limit + 1) - if split_at <= 0: + # A split which leaves nothing but blanks on the current line -- + # the only space is the statement's own indentation, as for the + # space-free Kleiss-Kuijf JAMPF lines -- emits an empty physical + # line, and the continuation which follows it is then attached to + # the *previous* statement. Break mid-token instead. + if split_at <= 0 or not remaining[:split_at+1].strip(): split_line.append(remaining[:char_limit]) remaining = remaining[char_limit:] else: diff --git a/tests/unit_tests/madevent/test_hel_recycle.py b/tests/unit_tests/madevent/test_hel_recycle.py new file mode 100644 index 0000000000..06d8edc59a --- /dev/null +++ b/tests/unit_tests/madevent/test_hel_recycle.py @@ -0,0 +1,81 @@ +############################################################################## +# +# Copyright (c) 2010 The MadGraph Development team and Contributors +# +# This file is a part of the MadGraph 5 project, an application which +# automatically generates Feynman diagrams and matrix elements for arbitrary +# high-energy processes in the Standard Model and beyond. +# +# It is subject to the MadGraph license which should accompany this +# distribution. +# +# For more information, please visit: http://madgraph.phys.ucl.ac.be +# +################################################################################ +""" Fixed-form line wrapping of the helicity recycled matrix element """ + +from __future__ import absolute_import +import unittest + +import madgraph.madevent.hel_recycle as hel_recycle + + +class TestDoMultiline(unittest.TestCase): + """do_multiline breaks a statement over fixed-form continuation lines. + + A continuation line is attached to whatever statement precedes it, so a + physical line which holds nothing but blanks must never be emitted: the + continuation after it would silently extend the previous statement.""" + + # the Kleiss-Kuijf color flow JAMPs of g g > g g g on the DDM basis: long + # enough to wrap and with no space of their own to wrap at, so the only + # candidate split point is inside the statement's indentation + JAMPF = ' JAMPF(2,1)=+2D0*(-IMAG1*JAMP(3,1)-IMAG1*JAMP(4,1)' \ + '-IMAG1*JAMP(6,1))' + + def physical_lines(self, line): + return hel_recycle.do_multiline(line).split('\n') + + def assertWrapIsValid(self, line): + lines = self.physical_lines(line) + for i, physical in enumerate(lines): + self.assertLessEqual(len(physical), 132, + 'physical line %d is too long' % i) + if i: + self.assertTrue(physical.lstrip().startswith('$'), + 'continuation line %d lost its marker' % i) + for i, physical in enumerate(lines[:-1]): + self.assertTrue(physical.strip(), + 'physical line %d is blank, so the continuation ' + 'which follows it joins the previous statement' % i) + # and nothing of the statement is lost on the way + rebuilt = ''.join(p.lstrip()[1:] if i else p + for i, p in enumerate(lines)) + self.assertEqual(rebuilt.replace(' ', ''), line.replace(' ', '')) + + def test_no_blank_physical_line_without_a_space_to_wrap_at(self): + """A statement whose only space is its indentation wraps mid-token""" + + self.assertWrapIsValid(self.JAMPF) + self.assertEqual(len(self.physical_lines(self.JAMPF)), 2) + + def test_short_line_is_untouched(self): + """Nothing happens below the limit""" + + short = ' JAMPF(1,1)=+2D0*(+IMAG1*JAMP(6,1))' + self.assertEqual(hel_recycle.do_multiline(short), short) + + def test_wraps_at_a_space_when_there_is_one(self): + """The usual case still breaks at the last space that fits""" + + line = ' JAMP(2,1) = (-1.000000000000000D+00)*AMP( K,8)+' \ + '(-1.000000000000000D+00)*AMP( K,11)+(-1.0D+00)*TMP_JAMP(20)' + self.assertWrapIsValid(line) + self.assertTrue(self.physical_lines(line)[0].endswith(' ')) + + def test_every_wrap_width_around_the_limit_is_valid(self): + """Sweep the statement length across the wrap width""" + + for pad in range(40): + line = ' JAMPF(2,1)=+2D0*(' + 'X' * pad + ')' + self.assertWrapIsValid(line)