diff --git a/madgraph/core/color_algebra.py b/madgraph/core/color_algebra.py index 8b3ac0db06..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 #=============================================================================== @@ -1128,18 +1135,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..0437ab2684 100755 --- a/madgraph/core/color_amp.py +++ b/madgraph/core/color_amp.py @@ -21,6 +21,9 @@ import collections import copy import fractions +import itertools +import itertools +import logging import operator import re import array @@ -35,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 #=============================================================================== @@ -53,6 +292,22 @@ 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']) + + # 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 @@ -237,8 +492,188 @@ 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 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.""" @@ -283,8 +718,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() @@ -322,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" @@ -346,6 +807,19 @@ 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 = {} + + # 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), \ @@ -448,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 @@ -529,113 +1008,698 @@ 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) + + +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. + + 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) + + 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 #=============================================================================== +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 = {} + _ddm_expansions = None 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 = [] + # 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 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.""" + + 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()) + + 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 + + if getattr(self._col_basis1, '_ddm_ends', None) and \ + getattr(self._col_basis2, '_ddm_ends', None): + self.setup_ddm_entries() 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: + 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._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 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 - # Fix indices in struct2 knowing summed indices in struct1 - # to avoid duplicates - new_struct2 = self.fix_summed_indices(struct1, struct2) + 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)) - # Build a canonical representation of the two immutable struct - canonical_entry, dummy = \ - color_algebra.ColorString().to_canonical(struct1 + \ - new_struct2) + 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 - 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)] - if is_symmetric: - self.inverted_col_matrix[result_fixed_Nc] = [(i2, i1)] 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) @@ -690,25 +1754,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/madgraph/interface/madgraph_interface.py b/madgraph/interface/madgraph_interface.py index d7e1ef7fd0..9de0da38aa 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 @@ -522,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).") @@ -2693,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" @@ -3155,8 +3157,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 +3242,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 +4244,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 +9363,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 +9828,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_cpp.py b/madgraph/iolibs/export_cpp.py index 354815b369..7edb51aaf0 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""" @@ -2688,17 +2701,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_mg7.py b/madgraph/iolibs/export_mg7.py index a6cc4f0518..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( @@ -205,8 +215,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_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/export_v4.py b/madgraph/iolibs/export_v4.py index 1fb735ea1b..900e254c36 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 @@ -41,12 +42,14 @@ 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 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 @@ -107,6 +110,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 @@ -211,10 +221,180 @@ 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, + jamp_optimiser.JampOptimiser): """Class to take care of exporting a set of matrix elements to Fortran (v4) format.""" @@ -224,7 +404,38 @@ class ProcessExporterFortran(VirtualExporter): 'output_options':{} } grouped_mode = False - jamp_optim = 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_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 + # 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' + # 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 + # 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 run_card_class = None use_flavor_mask = True @@ -1731,8 +1942,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]: @@ -2057,6 +2268,126 @@ 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 + 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 + + 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. + # 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 + + 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, + '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': [place[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.""" @@ -2064,6 +2395,34 @@ 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}] + + 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()) @@ -2096,11 +2455,216 @@ def get_color_data_lines(self, matrix_element, n=128): return ret_list + @staticmethod + 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(%s),%s=%d,%d) /%s/" % \ + (name, var, var, 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, + suffix=''): + """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 = 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%s()" % (proc_prefix, suffix)] + 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%s/ %sCF,%sDENOM" % \ + (proc_prefix, suffix, 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/" % \ 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. + + 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 = [] + 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'] = \ + ' 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).""" @@ -2123,20 +2687,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. @@ -2304,7 +2869,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 @@ -2373,22 +2939,34 @@ 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) 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='', + 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 + 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. @@ -2402,7 +2980,9 @@ 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 = [] for i, coeff_list in enumerate(color_amplitudes): # It might happen that coeff_list is empty if this function was @@ -2437,11 +3017,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 - value = (1j if coefficient[2] else 1)* coefficient[0] * coefficient[1] * fractions.Fraction(3)**coefficient[3] - 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], @@ -2460,32 +3035,37 @@ 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 - for key in all_element: - all_element[key] = complex(all_element[key]) - new_mat, defs = self.optimise_jamp(all_element) + # 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 + # 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) + 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): @@ -2504,30 +3084,104 @@ 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 + # 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 \ + 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, nb_amp) 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)) + recipes = self.jamp_orbit_recipes(defs, nb_amp) + self.jamp_recipes = recipes + + if recipes: + 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") + 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.") + 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(' TMP_JAMP(%d) = %s - %s ! used %d times' % (i,amp1, amp2, nb)) + 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") + 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(" %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(" %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(" %s(NGRAPHS+ITMP) = %s(TMP_JAMP_A(ITMP))" + " + TMP_JAMP_F(TMP_JAMP_L(5*ILEV)+ITMP" + "-TMP_JAMP_L(5*ILEV-2))" + "*%s(TMP_JAMP_B(ITMP))" + % (buffer, buffer, buffer)) + 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: + 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 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: - name = "TMP_JAMP(%d)" % -var + 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] + 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: @@ -2544,99 +3198,745 @@ def format(frac): return res_list, len(defs) - 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) - 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) - """ - 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 - - 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 index in all_index: - j1,j2,R = 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)] #= 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 - - - - + _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 + + @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.""" + + 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 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 + 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"), + 'color_fold_index': "\n".join( + self.get_int_data_lines("COLREP", lines, var='ICF')), + 'color_fold_gather': " JFOLD(:,:) = JAMP(COLREP(:),:)", + 'color_fold_array': 'JFOLD'} + + 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'] + + # 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 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 self.jamp_i_power(factor) is None: + 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) + + # 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, + 'complex_factor': complex_factor} + + @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 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 + by_index = dict((one[0], one) for one in defs) + levels = self.jamp_definition_levels(defs) + order, bounds, nb_general = [], [], 0 + for level in levels: + group = [[], [], []] + for index in level: + 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, 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(%s),%s=%d,%d) /%s/" % + (name, var, var, 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 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), + " 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), + ] + 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), []), + var="IJMP") + 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') + + 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 + 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.""" + + 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", + " PARAMETER (NB_TMP_JAMP=%d)" % len(recipes['defs']), + " INTEGER TMP_JAMP_A(NB_TMP_JAMP), TMP_JAMP_B(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, + ] + + 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 or not recipes.get('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)", + " %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,TMP_JAMP_E,HVAL,HKEY" % \ + proc_prefix, + ] + + add = [ + " SUBROUTINE %sJAMP_ADD(A,B,F,M,SWAP)" % proc_prefix, + "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 + [ + " 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))*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", + " 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_E(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 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, JIMGE", + ] + 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", + "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", + " A = -JIMG((-A-1)*NB_PERM+P)", + " EA = JIMGE((-TMP_JAMP_A(J)-1)*NB_PERM+P)", + " ENDIF", + " 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", + " B = -JIMG((-B-1)*NB_PERM+P)", + " EB = JIMGE((-TMP_JAMP_B(J)-1)*NB_PERM+P)", + " ENDIF", + "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) = M", + " JIMGE((J-1)*NB_PERM+P) = MOD(EA+SWAP,4)", + " 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)", + " TMP_JAMP_F(I) = IPOW(TMP_JAMP_E(I))", + " ENDDO", + " END", + ] + return add + body + + def jamp_tables_allowed(self): + """Whether the definitions may be read from a table rather than written + 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 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: only the + templates which know how to write the definitions it produces.""" + + if not super().jamp_orbit_allowed(matrix_element): + 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 + 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, power_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 + # 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 factor that goes with it""" + + if column > 0: + return amp_action[place][column] + 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 + 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 self.jamp_i_power(ratio) is None: + 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, self.jamp_i_power(ratio))) + + 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_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 + 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 + # 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) + + def get_pdf_lines(self, matrix_element, ninitial, subproc_group = False, vector=False): """Generate the PDF lines for the auto_dsig.f file""" @@ -3350,7 +4650,14 @@ 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 + # 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 @@ -3411,14 +4718,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) @@ -3915,6 +5225,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() @@ -4245,7 +5556,30 @@ 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'] =\ @@ -4254,6 +5588,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 @@ -4264,11 +5601,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. @@ -4311,7 +5651,103 @@ 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 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. + 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'] + + # 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))) + # 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" + 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] + flow_decl + + 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] + + flow_lines + [ + " 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: + # 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'] = "" + replace_dict['blas_routine'] = "" matrix_template = self.matrix_template if self.opt['export_format']=='standalone_msP' : @@ -4522,6 +5958,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): @@ -4588,9 +6026,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 @@ -5473,6 +6913,11 @@ class ProcessExporterFortranME(ProcessExporterFortran): MadEvent format.""" matrix_file = "matrix_madevent_v4.inc" + 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 + jamp_gather = True done_warning_tchannel = False default_opt = {'clean': False, 'complex_mass':False, @@ -6202,10 +7647,21 @@ 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) 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 @@ -6252,11 +7708,22 @@ 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 + 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 + + # 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'] = """ @@ -7543,6 +9010,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/jamp_optimiser.py b/madgraph/iolibs/jamp_optimiser.py new file mode 100644 index 0000000000..2b985c9bd1 --- /dev/null +++ b/madgraph/iolibs/jamp_optimiser.py @@ -0,0 +1,793 @@ +################################################################################ +# +# 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. + +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, 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 + +import bisect +import collections +import fractions +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') + + +class JampOptimiser(object): + """The common sub-expression search over the JAMP coefficient matrix. + + 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 + # (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 + # 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.""" + + 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, + 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 + 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. + + 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 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]", + len(all_element)) + start_time = time.time() + else: + 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: + 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 + + 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/madgraph/iolibs/template_files/madmatrix/MatrixElementKernels.cc b/madgraph/iolibs/template_files/madmatrix/MatrixElementKernels.cc index 872e4795e3..5aa8eae084 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 @@ -310,7 +311,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 @@ -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/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). diff --git a/madgraph/iolibs/template_files/madmatrix/color_sum.cc b/madgraph/iolibs/template_files/madmatrix/color_sum.cc index 30c6799932..903c6956a2 100644 --- a/madgraph/iolibs/template_files/madmatrix/color_sum.cc +++ b/madgraph/iolibs/template_files/madmatrix/color_sum.cc @@ -24,21 +24,21 @@ 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[ncolor * ncolor]; + T value[ncolorfold * ncolorfold]; }; - // The fptype2 version is the default used by kernels (supporting mixed floating point mode also in blas) - static __device__ fptype2 s_pNormalizedColorMatrix2[ncolor * ncolor]; + // The fptype2 version is the default used by kernels (supporting mixed floating point mode) + static __device__ fptype2 s_pNormalizedColorMatrixFold2[ncolorfold * ncolorfold]; #endif //-------------------------------------------------------------------------- @@ -51,7 +51,7 @@ 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 ) ); } } #endif @@ -71,64 +71,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 @@ -151,7 +149,7 @@ namespace mg5amcCpu #endif } #endif - +%(cpp_blas_color_sum)s //-------------------------------------------------------------------------- #ifdef MGONGPUCPP_GPUIMPL @@ -167,40 +165,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; @@ -216,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 @@ -261,60 +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; - // Get the address associated with the normalized color matrix in device memory + // 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 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) @@ -334,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 @@ -415,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..cbaef72dea 100644 --- a/madgraph/iolibs/template_files/madmatrix/color_sum.h +++ b/madgraph/iolibs/template_files/madmatrix/color_sum.h @@ -14,6 +14,11 @@ #include "CPPProcess.h" #include "GpuAbstraction.h" +#include +#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 @@ -22,6 +27,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 { @@ -76,6 +117,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..0c6083f4a0 --- /dev/null +++ b/madgraph/iolibs/template_files/madmatrix/color_sum_blas.inc @@ -0,0 +1,170 @@ +// 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 ifold = 0; ifold < ncolorfold; ifold++ ) + for( int jfold = 0; jfold < ncolorfold; jfold++ ) + value[ifold * ncolorfold + jfold] = + ( colorMatrix[ifold][jfold] / colorDenom[ifold] + colorMatrix[jfold][ifold] / colorDenom[jfold] ) / 2; + } + fptype2 value[ncolorfold * ncolorfold]; + }; + + // 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 ncolorfold 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 ncolorfold=60 and ncomb=128 this is a few hundred kB. + static thread_local std::vector scratch; + const size_t need = 4 * (size_t)ncolorfold * ncol + ncol; + if( scratch.size() < need ) scratch.resize( need ); + fptype2* JR = scratch.data(); + fptype2* JI = JR + (size_t)ncolorfold * ncol; + fptype2* ZR = JI + (size_t)ncolorfold * ncol; + fptype2* ZI = ZR + (size_t)ncolorfold * ncol; + fptype2* MEcol = ZI + (size_t)ncolorfold * 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 ) * ncolorfold; + // Gather the color flows the sum runs over: one per reversal pair when + // the basis folds, every flow otherwise (identity map). The folded + // matrix already carries what the dropped partner contributes. + for( int ifold = 0; ifold < ncolorfold; ifold++ ) + { + const int icol = colorFoldRep[ifold]; +#ifdef MGONGPU_CPPSIMD + JR[off + ifold] = cxreal( jamp_sv[ip * ncolor + icol] )[ieppV]; + JI[off + ifold] = cximag( jamp_sv[ip * ncolor + icol] )[ieppV]; +#else + JR[off + ifold] = cxreal( jamp_sv[ip * ncolor + icol] ); + JI[off + ifold] = cximag( jamp_sv[ip * ncolor + icol] ); +#endif + } + } + } + // Ztemp[ncolorfold][ncol] = ColorMatrix[ncolorfold][ncolorfold] * Jamps[ncolorfold][ncol], real and imaginary parts apart + blas_symm( ncolorfold, ncol, cfsym.value, JR, ZR ); + blas_symm( ncolorfold, 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 * ncolorfold; + fptype2 me = 0; + for( int ifold = 0; ifold < ncolorfold; ifold++ ) + me += JR[off + ifold] * ZR[off + ifold] + JI[off + ifold] * ZI[off + ifold]; + 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_class.inc b/madgraph/iolibs/template_files/madmatrix/process_class.inc index 59a6d07333..cbf55bacdb 100644 --- a/madgraph/iolibs/template_files/madmatrix/process_class.inc +++ b/madgraph/iolibs/template_files/madmatrix/process_class.inc @@ -56,6 +56,10 @@ 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 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/madgraph/iolibs/template_files/madmatrix/process_function_definitions.inc b/madgraph/iolibs/template_files/madmatrix/process_function_definitions.inc index 082c373aa2..1ae25f9b05 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,10 +812,10 @@ 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++ ) + 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! @@ -829,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; @@ -841,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/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..5a3a3f47b5 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,13 +110,13 @@ 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) #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 ) { @@ -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/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 b1523fc0e2..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 @@ -296,8 +296,11 @@ 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) + INTEGER NCOLOR_FLOW + PARAMETER (NCOLOR_FLOW=%(ncolor_flow)d) REAL*8 ZERO PARAMETER (ZERO=0D0) COMPLEX*16 IMAG1 @@ -321,10 +324,15 @@ C C LOCAL VARIABLES C INTEGER I,J,M,N - COMPLEX*16 ZTEMP, TMP_JAMP(%(nb_temp_jamp)i) - INTEGER CF(NCOLOR*(NCOLOR+1)/2) + COMPLEX*16 ZTEMP +%(jamp_tmp_decl)s +%(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 +%(jampflow_decl)s type(aloha) W(NWAVEFUNCS) C Needed for v4 models COMPLEX*16 DUM0,DUM1 @@ -364,11 +372,13 @@ C C COLOR DATA C %(color_data_lines)s +%(color_fold_index)s C ---------- 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 @@ -384,6 +394,7 @@ AMP(:) = (0d0,0d0) JAMP(:,:) = (0d0,0d0) %(jamp_lines)s +%(jampflow_lines)s if(init_mode)then DO I=1, NGRAPHS @@ -393,18 +404,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 @@ -415,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 @@ -465,4 +477,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 45bcdb8211..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 @@ -207,8 +207,11 @@ 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) + INTEGER NCOLOR_FLOW + PARAMETER (NCOLOR_FLOW=%(ncolor_flow)d) REAL*8 ZERO PARAMETER (ZERO=0D0) COMPLEX*16 IMAG1 @@ -234,11 +237,16 @@ 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 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 +%(jampflow_decl)s type(aloha) W(NWAVEFUNCS) C Needed for v4 models COMPLEX*16 DUM0,DUM1 @@ -272,11 +280,13 @@ C C COLOR DATA C %(color_data_lines)s +%(color_fold_index)s C ---------- 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. @@ -288,18 +298,20 @@ ${helas_calls} JAMP(:,:) = (0d0,0d0) DO K = 1, NCOMB ${jamp_lines} +%(color_fold_gather)s +%(jampflow_lines)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 @@ -308,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 @@ -330,4 +342,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 7b1d6bf1b9..ae44a888f5 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 @@ -271,10 +272,14 @@ C C LOCAL VARIABLES C INTEGER I,J,M,N - COMPLEX*16 ZTEMP, TMP_JAMP(%(nb_temp_jamp)i) - INTEGER CF(NCOLOR*(NCOLOR+1)) + COMPLEX*16 ZTEMP +%(jamp_tmp_decl)s +%(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) C Needed for v4 models COMPLEX*16 DUM0,DUM1 @@ -302,11 +307,13 @@ C C COLOR DATA C %(color_data_lines)s +%(color_fold_index)s C ---------- C BEGIN CODE C ---------- if (first) then first=.false. + CALL %(proc_prefix)sINIT_CF%(proc_id)s() %(fake_width_definitions)s endif @@ -317,18 +324,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 @@ -349,4 +357,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 ada2fd0eb1..5c512ee9ac 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 (%(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 @@ -228,7 +230,9 @@ 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 +%(jampflow_decl)s type(aloha) W(NWAVEFUNCS) COMPLEX*16 DUM0,DUM1 DATA DUM0, DUM1/(0d0, 0d0), (1d0, 0d0)/ @@ -248,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 @@ -344,11 +349,15 @@ 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 +%(jampflow_routine)s + SUBROUTINE %(proc_prefix)sGET_MATRIX(JAMP,MATRIX) C %(process_lines)s @@ -361,43 +370,75 @@ 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 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) + 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), TMP_JAMP(%(nb_temp_jamp)i) + 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)/ C 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 - 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) +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 + %(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)*JFOLD(I+J) ENDDO - MATRIX = MATRIX+ZTEMP*DCONJG(JAMP(I))/%(proc_prefix)sDENOM + CF_INDEX = CF_INDEX + NJ + MATRIX = MATRIX+ZTEMP*DCONJG(JFOLD(I))/%(proc_prefix)sDENOM ENDDO END +%(color_init_routine)s + +%(jamp_init_routine)s + +%(blas_routine)s SUBROUTINE %(proc_prefix)sGET_INTER(JAMP_1,JAMP_2, INTER) CF2PY INTENT(OUT) :: INTER @@ -416,6 +457,7 @@ CF2PY INTENT(IN) :: JAMP_2 C COLOR DATA C + CALL %(proc_prefix)sINIT_CF() INTER = (0.D0,0.D0) CF_INDEX = 0 @@ -547,7 +589,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/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/madmatrix/model_handling.py b/madmatrix/model_handling.py index b4d1bac787..320c8a53c4 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 @@ -1462,7 +1463,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 +1472,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 = '.' @@ -1483,6 +1488,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)" @@ -1503,6 +1510,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? @@ -1512,6 +1522,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 @@ -1729,6 +1740,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 @@ -1769,7 +1783,17 @@ 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'] = '' + 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 @@ -1904,8 +1928,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++ === @@ -1984,6 +2016,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 +2042,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() @@ -2058,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']) @@ -2127,33 +2186,158 @@ 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 - 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 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 + """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): + """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'] + 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 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 = 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)) # AV - replace the export_cpp.OneProcessExporterCPP method (improve formatting) def get_initProc_lines(self, matrix_element, color_amplitudes): @@ -2234,13 +2418,22 @@ 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 + # 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] @@ -2249,8 +2442,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 +2598,166 @@ 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, 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 + 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, + matrix_element=matrix_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 +2766,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 +2777,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(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 me = matrix_element.get('diagrams') matrix_element.reuse_outdated_wavefunctions(me) ###misc.sprint(multi_channel_map) @@ -2563,27 +2924,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..af6f988425 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] @@ -159,6 +167,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)') @@ -177,9 +198,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 @@ -230,10 +253,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 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 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 9775c450fc..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 @@ -765,3 +766,365 @@ 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) + + 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]) + + +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) 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) diff --git a/tests/unit_tests/iolibs/test_export_v4.py b/tests/unit_tests/iolibs/test_export_v4.py index f2d500e560..34e2db413c 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)) @@ -3796,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): @@ -10360,3 +10366,379 @@ 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) + + #=========================================================================== + # 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_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 + 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) + # 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) + 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) + # 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 + # 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)) 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)