diff --git a/MadSpin/decay.py b/MadSpin/decay.py index ca9ab1edc..be5772cbd 100755 --- a/MadSpin/decay.py +++ b/MadSpin/decay.py @@ -5427,6 +5427,9 @@ class DensityMatrix: # Cache diagonal masks by basis_id (depends only on helicities) _diag_cache = {} + # Same, as integer positions, for trace() + _diag_pos_cache = {} + # Cache tensor-product helicity tables by basis_id _tp_hel_cache = {} @@ -5629,6 +5632,21 @@ def _get_diag_mask_cached(self): DensityMatrix._diag_cache[self._basis_id] = mask return mask + def _get_diag_positions_cached(self): + """Positions of the diagonal entries, as a plain tuple of ints. + + trace() sums a handful of entries -- two, for a single decaying + fermion -- and numpy costs about 0.9 us to do that however few there + are, nearly all of it dispatch. Indexing the cached positions directly + is 8x faster at that size. Cached per basis_id alongside the mask. + """ + cached = DensityMatrix._diag_pos_cache.get(self._basis_id) + if cached is not None: + return cached + positions = tuple(int(i) for i in np.flatnonzero(self._diag_mask)) + DensityMatrix._diag_pos_cache[self._basis_id] = positions + return positions + # ------------------------------------------------------------------------- # Cached permutation for alignment by helicity labels # ------------------------------------------------------------------------- @@ -5678,7 +5696,11 @@ def scalar_multiplication(self, other): # Fastest correct path for map-built matrices if (self.map_density_matrix_ind is not None and self.map_density_matrix_ind is other.map_density_matrix_ind): - return np.sum(self.values * other.values) + # np.dot rather than np.sum(a*b): identical for complex (dot does + # not conjugate) but one call instead of two, and it skips the + # temporary the multiply would allocate. 0.23 us against 0.90 us + # on the 16-entry matrices this sees. + return np.dot(self.values, other.values) # Align by cached ordering for each basis self._ensure_sorted_view() @@ -5686,7 +5708,7 @@ def scalar_multiplication(self, other): a = self._sort_order b = other._sort_order - return np.sum(self.values[a] * other.values[b]) + return np.dot(self.values[a], other.values[b]) def tensor_product(self, other): """ @@ -5736,7 +5758,17 @@ def trace(self): """ Order-independent trace. """ - return np.sum(self.values[self._diag_mask]) + # Sum the diagonal entries by position. See + # _get_diag_positions_cached: at these sizes numpy's dispatch dwarfs + # the addition itself. + values = self.values + positions = self._get_diag_positions_cached() + if not positions: + return np.complex64(0) + total = values[positions[0]] + for i in positions[1:]: + total = total + values[i] + return total def print_full_matrix(self, precision=6): diff --git a/MadSpin/interface_madspin.py b/MadSpin/interface_madspin.py index 426251540..2234209b8 100755 --- a/MadSpin/interface_madspin.py +++ b/MadSpin/interface_madspin.py @@ -16,12 +16,15 @@ from __future__ import division from __future__ import absolute_import import collections +import contextlib +import json import logging import math import os import random import re import shutil +import subprocess import sys import time import glob @@ -81,6 +84,8 @@ def default_setup(self): self.add_param('density_tolerance', 1E-4, comment='Tolerance for deviation between density and full ME') self.add_param('decay_event_mult', 1E0, comment='Produce more events than needed so that MadSpin does not have to regenerate decay events') self.add_param('density_keep_jacobian', False, comment='keep track of the phase-space volume change related to the offshell reshuffling') + self.add_param('decay_generator', 'mg7', allowed=['mg7', 'madevent'], + comment='which backend generates the decay-event pools: mg7 (madmatrix/madspace) or the legacy Fortran madevent') ############################################################################ ## Special post-processing of the options ## @@ -135,6 +140,102 @@ def post_identical_particle_in_prod_and_decay(self, value, change_userdefine, ra if value not in ["crash", 'average', 'max', 'first']: raise Exception("value %s not supported for this parameter identical_in_prod_and_decay") +# Per-particle columns of mg7's "lhe_npy" event file, in the order they map +# onto lhe_parser.Particle's attributes below. +_NPY_PARTICLE_FIELDS = ('pdg_id', 'status_code', 'mother1', 'mother2', + 'color', 'anti_color', 'px', 'py', 'pz', 'energy', + 'mass', 'lifetime', 'spin') +_NPY_PARTICLE_ATTRS = ('pid', 'status', 'mother1', 'mother2', + 'color1', 'color2', 'px', 'py', 'pz', 'E', + 'mass', 'vtim', 'helicity') + + +class NpyDecayPool(object): + """A pool of decay events backed by mg7's numpy event file. + + mg7 can write its events either as LHE text or as a numpy structured array + carrying the same fields. MadSpin consumes a decay event per accept/reject + trial, so the text round trip is pure overhead: 19.7 us per event through + the LHE parser against 11.3 us building the same objects from the array. + + Exposes the part of lhe_parser.EventFile that the decay loop uses -- being + iterable, and carrying the channel's cross section -- so it drops into + evt_decayfile in place of an EventFile. + """ + + # Rows converted from numpy to python scalars per batch. Small enough that + # the converted rows stay a rounding error against the events they build, + # large enough that the per-call overhead disappears. + CHUNK = 4096 + + def __init__(self, path, cross): + import numpy + + self.name = path + self.cross = cross + # Memory-mapped: a pool is tens of MB and is read once, front to back. + self._records = numpy.load(path, mmap_mode='r') + self._count = len(self._records) + self._index = 0 + self._chunk = [] + self._chunk_index = 0 + + names = self._records.dtype.names + position = {name: i for i, name in enumerate(names)} + self._event_cols = [position[name] for name in + ('weight', 'scale', 'alpha_qed', 'alpha_qcd')] + self._particle_cols = [] + for index in range(1, len(names)): + if 'part%d_pdg_id' % index not in position: + break + self._particle_cols.append( + [position['part%d_%s' % (index, field)] + for field in _NPY_PARTICLE_FIELDS]) + + def __len__(self): + return self._count + + def __iter__(self): + return self + + def __next__(self): + if self._chunk_index >= len(self._chunk): + if self._index >= self._count: + raise StopIteration + # One vectorised conversion per chunk. Reading numpy scalars one + # field at a time instead costs more than parsing the LHE text did. + end = min(self._index + self.CHUNK, self._count) + self._chunk = self._records[self._index:end].tolist() + self._chunk_index = 0 + self._index = end + row = self._chunk[self._chunk_index] + self._chunk_index += 1 + + event = lhe_parser.Event() + event.wgt, event.scale, event.aqed, event.aqcd = ( + row[i] for i in self._event_cols) + for columns in self._particle_cols: + if row[columns[0]] == 0: + # Events shorter than the longest one in the file are padded + # with empty particles at the end. + break + particle = lhe_parser.Particle(event=event) + for attr, column in zip(_NPY_PARTICLE_ATTRS, columns): + setattr(particle, attr, row[column]) + event.append(particle) + event.nexternal = len(event) + # Mothers arrive as 1-based indices; the rest of MadSpin expects them + # resolved to the particles themselves, exactly as the parser leaves + # them. + event.assign_mother() + return event + + next = __next__ + + def close(self): + pass + + class MadSpinInterface(extended_cmd.Cmd): """Basic interface for madspin""" @@ -192,12 +293,67 @@ def __init__(self, event_path=None, *completekey, **stdin): self.err_branching_ratio = 0 self.me_run_name = "" # Events diretory name where to stotre the events (used by madevent) not use internally self.all_iden = {} - + # Wall-clock accounting per run phase (see _phase / _log_phase_timings). + # Always collected: the granularity is one entry per phase, not per + # event, so the overhead is irrelevant next to the phases themselves. + self._phase_times = collections.defaultdict(float) + self._phase_counts = collections.defaultdict(int) + # Which bucket generate_events charges its wall time to. The main + # pre-generation pass charges 'decay_event_generation'; the mid-loop + # pool refills in get_decay_from_file retarget it to + # 'decay_event_refill' so the two are told apart. + self._gen_phase = 'decay_event_generation' + if event_path: logger.info("Extracting the banner ...") self.do_import(event_path) - - + + # ------------------------------------------------------------------ + # Phase timing + # ------------------------------------------------------------------ + def _add_phase(self, name, dt, count=1): + """Charge ``dt`` seconds (and ``count`` occurrences) to phase ``name``.""" + self._phase_times[name] += dt + self._phase_counts[name] += count + + def _add_count(self, name, count): + """Record a pure counter (no wall time), e.g. events or trials.""" + self._phase_counts[name] += count + + @contextlib.contextmanager + def _phase(self, name): + """Context manager charging its body's wall time to phase ``name``.""" + start = time.time() + try: + yield + finally: + self._add_phase(name, time.time() - start) + + def _log_phase_timings(self): + """Emit one machine-readable line with the per-phase wall times. + + Kept as a single JSON payload on a fixed prefix so benchmark drivers + can pull the whole split out of a log with one regex, without having + to track the wording of the individual human-readable timing lines. + """ + if not self._phase_times: + return + payload = { + 'seconds': {k: round(v, 4) for k, v in sorted(self._phase_times.items())}, + 'counts': {k: v for k, v in sorted(self._phase_counts.items())}, + } + logger.critical('MadSpin phase timings: %s', json.dumps(payload, sort_keys=True)) + + def _finish_run(self): + """End-of-run reporting: total wall time, the phase split, and the + optional LHE parser timers.""" + if getattr(self, '_run_start', None) is not None: + self._add_phase('total', time.time() - self._run_start) + self._run_start = None + self._log_phase_timings() + self._log_lhe_timers() + + def setup_for_pure_decay(self): """this is for spinmode=none -> simple decay We go here if they are no banner. @@ -715,9 +871,12 @@ def do_launch(self, line): """end of the configuration launched the code""" (options, args) = self.parse_launch(line) + self._run_start = time.time() + self._phase_times.clear() + self._phase_counts.clear() if getattr(lhe_parser, "_ENABLE_LHE_TIMERS", False): lhe_parser.reset_lhe_timers() - + if options.name: self.me_run_name = options.name # Only use by MG5aMC else: @@ -736,25 +895,25 @@ def do_launch(self, line): if spinmode in ["none"]: out = self.run_bridge(line) - self._log_lhe_timers() + self._finish_run() return out elif spinmode.startswith("onshell"): if spinmode == "onshell_v1": out = self.run_onshell(line) else: out = self.run_onshell(line, density_method=True) - self._log_lhe_timers() + self._finish_run() return out elif spinmode == "PA": out = self.run_onshell(line, density_method=True) - self._log_lhe_timers() + self._finish_run() return out elif spinmode == "madspin_v1": # legacy MadSpin / decay-chain path: fall through to decay_all_events below pass elif spinmode == "madspin": out = self.run_onshell(line, density_method=True) - self._log_lhe_timers() + self._finish_run() return out elif spinmode == "bridge": raise Exception("Bridge mode not available.") @@ -763,7 +922,7 @@ def do_launch(self, line): if self.options['ms_dir'] and os.path.exists(pjoin(self.options['ms_dir'], 'madspin.pkl')): out = self.run_from_pickle() - self._log_lhe_timers() + self._finish_run() return out @@ -813,12 +972,16 @@ def do_launch(self, line): time_me_generation = time.time() self.update_status('generating Madspin matrix element') - generate_all = madspin.decay_all_events(self, self.banner, self.events_file, + generate_all = madspin.decay_all_events(self, self.banner, self.events_file, self.options) - logger.critical(f"Time for ME: {time.time()-time_me_generation:.2f} sec") + time_me_generation = time.time() - time_me_generation + logger.critical(f"Time for ME: {time_me_generation:.2f} sec") + self._add_phase('me_generation', time_me_generation) self.update_status('running MadSpin') - generate_all.run() - + with self._phase('decay_loop'): + generate_all.run() + + self.branching_ratio = generate_all.branching_ratio self.cross = generate_all.cross self.error = generate_all.error @@ -863,7 +1026,7 @@ def do_launch(self, line): misc.call(['tar','-czpf','RunMaterial.tar.gz','RunMaterial'], cwd=run_dir) shutil.rmtree(pjoin(run_dir,'RunMaterial')) - self._log_lhe_timers() + self._finish_run() def run_from_pickle(self): import madgraph.iolibs.save_load_object as save_load_object @@ -1349,6 +1512,93 @@ def load_model(self, name, use_mg_default, complex_mass=False): self.mg5cmd._curr_model = self.model self.mg5cmd.process_model() + @property + def decay_generator(self): + """Backend used to produce the decay-event pools. + + Gridpack mode stays on Fortran madevent: it drives the decay directory + through run.sh, which the mg7 output does not provide. + """ + generator = self.options['decay_generator'] + if generator == 'mg7' and self.options['ms_dir']: + logger.info('gridpack mode: generating decay events with madevent ' + 'rather than mg7') + return 'madevent' + return generator + + def generate_events_mg7(self, decay_dir, nb_event): + """Run the mg7 (madmatrix/madspace) generator in ``decay_dir``. + + Returns ``(event_file, partial_width)``. The partial width is read back + from the LHE block, which for a 1 -> n process carries a width in + GeV rather than a cross section in pb. + """ + run_card_path = pjoin(decay_dir, 'Cards', 'run_card.toml') + run_card = banner.RunCardMG7(run_card_path) + run_card['generation']['events'] = int(nb_event) + # The pool is read straight back into MadSpin, so write it as the numpy + # event file rather than LHE text: same fields, no parser. Nothing here + # is showered or analysed, so no tool needs the LHE. + run_card['run']['output_format'] = 'lhe_npy' + run_card.write(run_card_path) + with open(pjoin(decay_dir, 'Cards', 'param_card.dat'), 'w') as fsock: + fsock.write(self.banner['slha']) + + events_dir = pjoin(decay_dir, 'Events') + before = set(misc.glob('*', events_dir)) if os.path.isdir(events_dir) else set() + # Run out of process: the launcher chdirs, installs signal handlers and + # holds its own madspace context, none of which should land in MadSpin's + # interpreter next to the f2py matrix elements. + # Append rather than truncate: a pool refill reruns the launcher in the + # same directory, and overwriting would throw away the log of the run + # that did the matrix-element compile. + log_path = pjoin(decay_dir, 'mg7_generation.log') + with open(log_path, 'a') as logfile, self._phase('decay_mg7_launch'): + returncode = misc.call( + [sys.executable, pjoin(decay_dir, 'bin', 'generate_events'), '-f'], + cwd=decay_dir, stdout=logfile, stderr=subprocess.STDOUT) + if returncode: + raise self.InvalidCmd( + 'the mg7 decay generator failed in %s (exit %s); see %s' + % (decay_dir, returncode, log_path)) + after = set(misc.glob('*', events_dir)) if os.path.isdir(events_dir) else set() + new_runs = sorted(after - before) + if not new_runs: + raise self.InvalidCmd( + 'the mg7 decay generator produced no run directory in %s' % events_dir) + run_dir = new_runs[-1] + + events_path = pjoin(run_dir, 'events.npy') + if not os.path.exists(events_path): + raise self.InvalidCmd( + 'the mg7 decay generator produced no events.npy in %s' % run_dir) + + # The numpy event file carries no block, so the partial width and + # the timing breakdown both come from the run's info.json. The launcher + # reports how long the integration itself took; the rest of the + # subprocess wall time is interpreter start-up plus the one-off + # matrix-element compile, which is what dominates a decay directory. + info_path = pjoin(run_dir, 'info.json') + try: + with open(info_path) as fsock: + info = json.load(fsock) + except (OSError, ValueError) as error: + raise self.InvalidCmd( + 'could not read the mg7 result from %s: %s' % (info_path, error)) + try: + width = float(info['process']['mean']) + except (KeyError, TypeError, ValueError) as error: + raise self.InvalidCmd( + 'no partial width in %s: %s' % (info_path, error)) + run_times = info.get('run_times', {}) + if isinstance(run_times, dict): + self._add_phase('decay_mg7_integrate', + sum(stage.get('wall_time_sec', 0.) + for stage in run_times.values() + if isinstance(stage, dict))) + + return NpyDecayPool(events_path, width), width + def generate_events(self, pdg, nb_event, mg5, restrict_file=None, cumul=False, output_width=False): """generate new events for this particle @@ -1379,24 +1629,27 @@ def generate_events(self, pdg, nb_event, mg5, restrict_file=None, cumul=False, if restrict_file and i not in restrict_file: continue decay_dir = pjoin(self.path_me, "decay_%s_%s" %(str(pdg).replace("-","x"),i)) + # Pin the output format explicitly rather than relying on MG5's + # default: the directory is driven below by the matching runner, so + # the two must agree. + output_format = self.decay_generator if not os.path.exists(decay_dir): - if cumul: - mg5.exec_cmd("generate %s" % proc) - for j,proc2 in enumerate(self.list_branches[name][1:]): - misc.sprint(proc2) - if restrict_file and j not in restrict_file: - raise Exception # Do not see how this can happen - mg5.exec_cmd("add process %s" % proc2) - # Force the Fortran madevent output: the decay directory is - # driven below through MadEventCmdShell, so it must have the - # madevent structure regardless of MG5's default output mode - # (which is 'mg7' in MadGraph7). - mg5.exec_cmd("output madevent %s -f" % decay_dir) - else: - misc.sprint(proc) - mg5.exec_cmd("generate %s" % proc) - mg5.exec_cmd("output madevent %s -f" % decay_dir) - + # Amplitude generation and code writing, paid once per decay + # directory (a pool refill reuses it). + with self._phase('decay_me_generate'): + if cumul: + mg5.exec_cmd("generate %s" % proc) + for j,proc2 in enumerate(self.list_branches[name][1:]): + misc.sprint(proc2) + if restrict_file and j not in restrict_file: + raise Exception # Do not see how this can happen + mg5.exec_cmd("add process %s" % proc2) + else: + misc.sprint(proc) + mg5.exec_cmd("generate %s" % proc) + with self._phase('decay_me_output'): + mg5.exec_cmd("output %s %s -f" % (output_format, decay_dir)) + options = dict(mg5.options) if self.options['ms_dir']: # we are in gridpack mode -> create it @@ -1447,6 +1700,15 @@ def generate_events(self, pdg, nb_event, mg5, restrict_file=None, cumul=False, misc.call(['tar', '-xzpvf', 'run_01_gridpack.tar.gz'], cwd=decay_dir,stdout=devnull, stderr=-2) devnull.close() # Now generate the events + if self.decay_generator == 'mg7': + # mg7 delivers exactly the requested number of events, so ask + # for what is needed rather than madevent's 0.8x undershoot. + out[i], pwidth = self.generate_events_mg7(decay_dir, nb_event) + if output_width: + width = width + pwidth if cumul else width * pwidth + if cumul: + break + continue if not self.options['ms_dir']: if decay_dir in self.me_int: me5_cmd = self.me_int[decay_dir] @@ -1519,6 +1781,11 @@ def generate_events(self, pdg, nb_event, mg5, restrict_file=None, cumul=False, break time_gen_dec = time.time()-time_gen_dec logger.critical(f"Time for decay event generation = {time_gen_dec:.1f} sec") + # Charge to whichever bucket the caller selected (pre-generation pass + # vs. mid-loop refill) and record how many decay events were asked for, + # so a benchmark can tell "slow generator" from "generated too many". + self._add_phase(self._gen_phase, time_gen_dec) + self._add_count('%s_events_requested' % self._gen_phase, nb_event) if not output_width: return out else: @@ -1565,8 +1832,9 @@ def run_onshell(self, line, density_method=False): shutil.rmtree(name) self.events_file.close() - if self.events_file.name.endswith('.gz'): - misc.gunzip(self.events_file.name) + # Read the input where it is: EventFile handles a gzipped file itself. + # Unpacking it here only to repack identical content at the end of the + # run cost 6 s on a 100k-event sample, and the file is never written to. orig_lhe = lhe_parser.EventFile(self.events_file.name) if self.options['fixed_order']: orig_lhe.eventgroup = True @@ -1767,14 +2035,33 @@ def run_onshell(self, line, density_method=False): self.all_matrix = {} time_me_generation = time.time() - time_me_generation logger.critical(f"Time ME generation: {time_me_generation:.2f} sec") + self._add_phase('me_generation', time_me_generation) #4. determine the maxwgt + # This phase is not free: it probes Nevents_for_max_weight production + # events with max_weight_ps_point decay trials each (75 x 400 = 30000 + # trials by default), each consuming decay events and a density ME call. + time_maxwgt = time.time() maxwgt = self.get_maxwgt_for_onshell(orig_lhe, evt_decayfile, decay_dict) + time_maxwgt = time.time() - time_maxwgt + logger.critical(f"Time for max-weight scan = {time_maxwgt:.2f} sec") + # The refills triggered from inside the scan were charged to + # decay_event_refill by generate_events; keep max_weight as the + # inclusive number and let the benchmark subtract if it wants to. + self._add_phase('max_weight_scan', time_maxwgt) #5. generate the decay (for each production event) orig_lhe.seek(0) - output_lhe = lhe_parser.EventFile(orig_lhe.name.replace('.lhe', '_decayed.lhe'), 'w') + # Name the output after the input with any .gz stripped: now that the + # input is read in place, orig_lhe.name may carry it, and appending + # _decayed to that would ask EventFile to write a gzip stream here + # instead of the plain file the gzip step below expects. + input_base = orig_lhe.name + if input_base.endswith('.gz'): + input_base = input_base[:-3] + output_lhe = lhe_parser.EventFile( + input_base.replace('.lhe', '_decayed.lhe'), 'w') if self.options['fixed_order']: output_lhe.eventgroup = True @@ -1889,10 +2176,17 @@ def run_onshell(self, line, density_method=False): output_lhe.write_events(full_evt) output_lhe.write('\n') + # The accept/reject loop proper (decay-pool reads, density MEs, + # reshuffling, event writing). Refills triggered from inside it were + # charged separately by generate_events. + self._add_phase('decay_loop', time.time() - start) # Log unweighting efficiency (can be turned off) n_processed = curr_event + 1 n_written = n_processed - nb_loose_skip eff = float(n_written) / nb_try if nb_try else 0.0 + self._add_count('unweighting_trials', nb_try) + self._add_count('events_written', n_written) + self._add_count('events_processed', n_processed) logger.critical( "MadSpin unweight efficiency: %.4f (%d written / %d trials, %.2f trials/event)", eff, n_written, nb_try, (1.0 / eff if eff else float("inf")) @@ -1917,28 +2211,31 @@ def run_onshell(self, line, density_method=False): self.efficiency = br_correction else: self.efficiency = 1 # to let me5 to write the correct number of events - # Re-gzip the input events file (gunzipped at the start of this - # routine) and the decayed output, matching the legacy MadSpin path - # so downstream code (banners, crossx.html) finds the *.lhe.gz files - # it expects. - try: - output_lhe.close() - except Exception: - pass - try: - input_evt_path = self.events_file.name - if input_evt_path.endswith('.lhe') and os.path.exists(input_evt_path): - misc.gzip(input_evt_path) - except Exception as exc: - logger.warning('Could not re-gzip MadSpin input file %s: %s', - getattr(self.events_file, 'name', '?'), exc) - try: - decayed_path = output_lhe.name - if decayed_path.endswith('.lhe') and os.path.exists(decayed_path): - misc.gzip(decayed_path) - except Exception as exc: - logger.warning('Could not gzip MadSpin decayed output %s: %s', - output_lhe.name, exc) + # Gzip the decayed output, matching the legacy MadSpin path so + # downstream code (banners, crossx.html) finds the *.lhe.gz file it + # expects. The input is left exactly as it was found: it is read in + # place, gzipped or not, and never modified. + with self._phase('output_gzip'): + try: + output_lhe.close() + except Exception: + pass + try: + input_evt_path = self.events_file.name + if input_evt_path.endswith('.lhe') and os.path.exists(input_evt_path): + # Only reachable when the caller handed us a plain .lhe; + # downstream still expects to find it gzipped. + misc.gzip(input_evt_path) + except Exception as exc: + logger.warning('Could not gzip MadSpin input file %s: %s', + getattr(self.events_file, 'name', '?'), exc) + try: + decayed_path = output_lhe.name + if decayed_path.endswith('.lhe') and os.path.exists(decayed_path): + misc.gzip(decayed_path) + except Exception as exc: + logger.warning('Could not gzip MadSpin decayed output %s: %s', + output_lhe.name, exc) logger.info('Done so far. output written in %s' % output_lhe.name) logger.critical(f"Time for decay = {time.time()-start:.2f} sec") @@ -2064,8 +2361,12 @@ def get_decay_from_file(self,production, evt_decayfile, nb_remain): burn = max(1.0, float(same_pdg) / float(nb_decay)) needed = int(math.ceil(1.10 * burn * nb_remain / eff)) needed = min(200000, max(needed, 1000)) - with misc.MuteLogger(["madgraph", "madevent", "ALOHA", "cmdprint"], [50,50,50,50]): - new_file = self.generate_events(particle.pdg, needed, self.mg5cmd, [decay_file_nb]) + self._gen_phase = 'decay_event_refill' + try: + with misc.MuteLogger(["madgraph", "madevent", "ALOHA", "cmdprint"], [50,50,50,50]): + new_file = self.generate_events(particle.pdg, needed, self.mg5cmd, [decay_file_nb]) + finally: + self._gen_phase = 'decay_event_generation' evt_decayfile[particle.pdg].update(new_file) decay_file = evt_decayfile[particle.pdg][decay_file_nb] continue @@ -2074,6 +2375,83 @@ def get_decay_from_file(self,production, evt_decayfile, nb_remain): return out + def batch_decay_densities(self, production, trials): + """boost and evaluate the decay densities of a set of trials at once. + + trials is a list of decays dictionaries, all drawn against the same + production event. Returns one list of DensityMatrix per trial, in the + traversal order calculate_matrix_element_from_density walks + (for pdg in decays: for i in range(N)), ready to be handed back to it as + decay_densities -- or None when the batched path does not apply, in + which case the caller must keep the per-trial path. + + The boost into the lab frame mutates the decay event in place, so it is + done here exactly once per event; passing the result to + calculate_matrix_element_from_density is what stops it boosting again. + """ + if not trials: + return [] + if self.generate_all.mode != 'density': + return None + if self.options['spinmode'] not in ('PA', 'onshell'): + # Outside the pole approximation the production and every decay are + # reshuffled *inside* calculate_matrix_element_from_density, so the + # momenta a density would need are not known before it runs. + return None + prod_static = getattr(production, '_ms_density_static', None) + if not prod_static: + # Populated by the first trial through the unbatched path; without + # it we do not know what to boost by or which helicities to ask for. + return None + + decays_key = prod_static['decays_key'] + init_part = prod_static['init_part'] + helicities = prod_static['helicities'] + nchanging = prod_static['nchanging'] + + # Flatten each trial into the traversal order, and check it presents the + # decay structure prod_static was built for -- the cache is only reused + # while decays_key holds, and the helicity slots are indexed by it. + flat = [] + for decays in trials: + if tuple(decays.keys()) != decays_key: + return None + events = [dec for pdg in decays_key for dec in decays[pdg]] + if len(events) != nchanging: + return None + flat.append(events) + + # One boost per decaying particle of the production event, shared by + # every trial: nothing moves the production between them here. + boosts = [] + for part in init_part: + boost = -1 * lhe_parser.FourMomentum(part) + boost.E *= -1 + boosts.append(boost) + for events in flat: + for slot, event in enumerate(events): + event.boost(boosts[slot]) + + # Slots that share a helicity structure share a fortran call: for + # t t~ both parents are spin 1/2, so the two slots of every trial go in + # together and the batch is 2 x the number of trials. + by_hel = {} + for slot in range(nchanging): + by_hel.setdefault(tuple(helicities[slot]), []).append(slot) + + out = [[None] * nchanging for _ in flat] + for slots in by_hel.values(): + allow_hel = helicities[slots[0]] + keys = [(t, slot) for t in range(len(flat)) for slot in slots] + densities = self.get_density_batch([flat[t][slot] for t, slot in keys], + position=[1], + allow_hel=allow_hel, + ncomb=len(allow_hel), + dimension=len(allow_hel)) + for (t, slot), density in zip(keys, densities): + out[t][slot] = density + return out + def get_maxwgt_for_onshell(self, orig_lhe, evt_decayfile, decay_dict): """determine the maximum weight for the onshell (or similar) strategy""" #print(f"decay_dict = {decay_dict} - length = {len(decay_dict)}") @@ -2109,17 +2487,38 @@ def get_maxwgt_for_onshell(self, orig_lhe, evt_decayfile, decay_dict): base_event = base_event[0] # Cache production density matrix density_matrix_prod = None + # The scan keeps every trial -- there is no accept/reject here -- so + # every decay set of this production event can be drawn up front and + # their densities evaluated in one batched fortran call rather than + # one call per trial. get_decay_from_file also refills the pools. + # Outside the pole approximation the decays are reshuffled *inside* + # the ME call, so their momenta are not known beforehand: there the + # draws stay lazy and the loop is exactly the one it always was. + npoints = self.options['max_weight_ps_point'] + can_batch = (self.generate_all.mode == 'density' and + self.options['spinmode'] in ('PA', 'onshell')) + all_decays = [self.get_decay_from_file(base_event, evt_decayfile, nevents-i) + for _ in range(npoints)] if can_batch else None + decay_densities = None # Loop over decays - for j in range(self.options['max_weight_ps_point']): - decays = self.get_decay_from_file(base_event, evt_decayfile, nevents-i) - #carefull base_event is modified by the following function + for j in range(npoints): + decays = all_decays[j] if can_batch else \ + self.get_decay_from_file(base_event, evt_decayfile, nevents-i) + #carefull base_event is modified by the following function if density_matrix_prod is None: _, wgt, density_matrix_prod = self.get_onshell_evt_and_wgt( base_event, decays, decay_dict, build_event=False) #print(f"wgt1 = {wgt}") + # This first trial populated production._ms_density_static, + # which says what to boost by and which helicities to ask + # for; the remaining trials then go in as one batch. + if can_batch: + decay_densities = self.batch_decay_densities( + base_event, all_decays[j+1:]) else: wgt = self.get_onshell_evt_and_wgt( - base_event, decays, decay_dict, density_matrix_prod, build_event=False)[1] + base_event, decays, decay_dict, density_matrix_prod, build_event=False, + decay_densities=decay_densities[j-1] if decay_densities else None)[1] #print(f"wgt2 = {wgt}") #print(f"Event {i} , PS point {j}, wgt for max = {wgt}") jac = 1 @@ -2160,11 +2559,15 @@ def get_maxwgt_for_onshell(self, orig_lhe, evt_decayfile, decay_dict): return base_max_weight - def get_onshell_evt_and_wgt(self, production, decays, decay_dict, prod_density_cached=None, build_event=True): + def get_onshell_evt_and_wgt(self, production, decays, decay_dict, prod_density_cached=None, + build_event=True, decay_densities=None): """ return the onshell wgt for the production event associated to the decays return also the full event with decay. Carefull this modifies production event (pass to the full one) - build_event: if False (density mode) compute weight without building event""" + build_event: if False (density mode) compute weight without building event + decay_densities: pre-computed decay density matrices (see + calculate_matrix_element_from_density); the decay events must + already have been boosted to the lab frame by the caller""" #print("\n\n\n\n\n======== debug get_onshell_evt_and_wgt =========") density_pole_approximation = self.options['spinmode'] in ['PA', 'onshell'] density_do_reshuffle = self.options['spinmode'] == 'PA' @@ -2214,14 +2617,17 @@ def get_onshell_evt_and_wgt(self, production, decays, decay_dict, prod_density_c nb_prod_final = sum(1 for p in production if int(p.status) == 1) if nb_prod_final > 1 and (not density_pole_approximation or density_do_reshuffle): + # decay_dict holds [width, mass, color, spin] per pdg, read from + # the same param card at the start of the run. Asking the banner + # again on every trial cost ~10% of the decay loop -- 4 lookups + # per trial, each walking the card -- for values that cannot + # change. + bw_cut = self.options['BW_cut'] + if bw_cut < 0: + bw_cut = 15 for pdg in decays: for dec in decays[pdg]: - pole = self.banner.get('param', 'mass', abs(pdg)).value - width = self.banner.get('param', 'decay', abs(pdg)).value - if self.options['BW_cut'] <0: - bw_cut = 15 - else: - bw_cut = self.options['BW_cut'] + width, pole = decay_dict[pdg][0], decay_dict[pdg][1] min_mass = pole - bw_cut * width max_mass = min(pole + bw_cut * width,full_dqrts) dec[0].new_mass = lhe_parser.Event.generate_random_mass(pole, width, min_mass, max_mass) @@ -2232,9 +2638,9 @@ def get_onshell_evt_and_wgt(self, production, decays, decay_dict, prod_density_c gap += math.atan((max_mass**2-pole**2)/pole*width) jac *= gap/math.pi if prod_density_cached is None: - full_me, prod_density_cached, prod_diag, dec_diag = self.calculate_matrix_element_from_density(production, decays, decay_dict) - else: - full_me, _, prod_diag, dec_diag = self.calculate_matrix_element_from_density(production, decays, decay_dict, prod_density_cached) + full_me, prod_density_cached, prod_diag, dec_diag = self.calculate_matrix_element_from_density(production, decays, decay_dict, decay_densities=decay_densities) + else: + full_me, _, prod_diag, dec_diag = self.calculate_matrix_element_from_density(production, decays, decay_dict, prod_density_cached, decay_densities=decay_densities) #print(f"full_me from density = {full_me}") full_event = None @@ -2291,8 +2697,18 @@ def get_onshell_evt_and_wgt(self, production, decays, decay_dict, prod_density_c return full_event, full_me/(production_me*decay_me)*jac, prod_density_cached - def calculate_matrix_element_from_density(self, production, decays, decay_dict, prod_density_cached=None): - """routine to return the matrix element from density matrices""" + def calculate_matrix_element_from_density(self, production, decays, decay_dict, + prod_density_cached=None, + decay_densities=None): + """routine to return the matrix element from density matrices + + decay_densities, when given, holds one DensityMatrix per decaying + particle in the traversal order below (for pdg in decays: for i in + range(N)), already evaluated in the lab frame. The loop then skips both + the boost and the fortran call: the caller has done them, in one batched + call for a whole set of trials. The caller owns the boost in that case + -- it mutates the decay event in place and must happen exactly once. + """ # ------------------------------------------------------------------ # Load f2py module and build pdg2prefix map if needed (unchanged logic) @@ -2515,19 +2931,23 @@ def _decay_signature(dec_evt): for i_decay_event in range(N): current_decay_event = decay_event_list[i_decay_event] - # boost to lab frame using corresponding production particle momentum - part = init_part[decaying_idx + i_decay_event] - boost = -1 * lhe_parser.FourMomentum(part) - boost.E *= -1 - current_decay_event.boost(boost) - - density_dec_tmp = self.get_density( - current_decay_event, - position=[1], - allow_hel=helicities[decaying_idx + i_decay_event], - ncomb=len(helicities[decaying_idx + i_decay_event]), - dimension=len(helicities[decaying_idx + i_decay_event]) - ) + if decay_densities is not None: + # Boosted and evaluated by the caller, in one batched call. + density_dec_tmp = decay_densities[decaying_idx + i_decay_event] + else: + # boost to lab frame using corresponding production particle momentum + part = init_part[decaying_idx + i_decay_event] + boost = -1 * lhe_parser.FourMomentum(part) + boost.E *= -1 + current_decay_event.boost(boost) + + density_dec_tmp = self.get_density( + current_decay_event, + position=[1], + allow_hel=helicities[decaying_idx + i_decay_event], + ncomb=len(helicities[decaying_idx + i_decay_event]), + dimension=len(helicities[decaying_idx + i_decay_event]) + ) if density_dec is None: density_dec = density_dec_tmp @@ -2647,13 +3067,96 @@ def get_density(self, event, position, allow_hel, ncomb, dimension): alphas=event.aqcd, scale2=event.scale**2) #print(f"density_array = {density_array}") - density_matrix = madspin.DensityMatrix(density_array, - n_changing, - allow_hel, + density_matrix = madspin.DensityMatrix(density_array, + n_changing, + allow_hel, dimension) return density_matrix - + + def get_density_batch(self, events, position, allow_hel, ncomb, dimension): + """get_density for a list of events sharing one helicity structure. + + The density itself is cheap -- 4 us of fortran against 8 us of python + preparing its arguments and 3 us building the DensityMatrix -- so the + win here is not the fortran loop but paying the f2py entry, the + argument marshalling and the numpy conversions once for the whole + list. POS and ALLOW_HEL describe the helicity structure and are shared + by construction; everything else is per point, because a decay pool + mixes flavours and nothing guarantees the points share a scale. + + Returns one DensityMatrix per event, in input order. + """ + import numpy as np + + n_changing = len(position) + if n_changing == 0: + raise ValueError("Error in get_density_batch: 'position' must contain at least one position index") + if len(allow_hel) % n_changing != 0: + raise ValueError("Error in get_density_batch: inconsistent 'allow_hel' and 'position' lengths") + + merged_map = self._revert_merged or None + merged_particles = self.model.get('merged_particles') or {} + + # Same derivation as get_density, per event. One fortran call carries a + # single NEXT, so the points are grouped by external multiplicity: a + # decay pool can hold 1 -> 2 and 1 -> 3 channels for the same parent. + momenta = [None] * len(events) + pdgs = [None] * len(events) + groups = {} + for k, event in enumerate(events): + orig_order = getattr(event, '_ms_orig_order_for_density', None) + if orig_order is None: + _, orig_order, _, _ = self.get_pdir(event) + event._ms_orig_order_for_density = orig_order + try: + p = event.get_momenta(orig_order, merged_map=merged_map) + except Exception: + # Safety fallback for unusual event structures. + all_p = event.get_all_momenta(orig_order, merged_map=merged_map) + assert len(all_p) == 1, "Error: get_density_batch can only be called for single phase-space points" + p = all_p[0] + pdg_template = list(orig_order[0]) + list(orig_order[1]) + need_raw_pdg = (self._revert_merged and + any(abs(pid) in merged_particles for pid in pdg_template)) + momenta[k] = p + pdgs[k] = event.get_pdg(p) if need_raw_pdg else pdg_template + groups.setdefault((len(p), len(pdgs[k])), []).append(k) + + out = [None] * len(events) + for (next_, npdg), idx in groups.items(): + nbatch = len(idx) + # (0:3, next, nbatch) fortran-ordered. get_momenta hands back one + # (E,px,py,pz) tuple per particle, so the transpose that + # invert_momenta walks in python per point is the array layout here. + P = np.empty((4, next_, nbatch), dtype=float, order='F') + all_pdgs = np.empty((npdg, nbatch), dtype=np.int32, order='F') + alphas = np.empty(nbatch, dtype=float) + scale2 = np.empty(nbatch, dtype=float) + for c, k in enumerate(idx): + P[:, :, c] = np.asarray(momenta[k], dtype=float).T + all_pdgs[:, c] = pdgs[k] + alphas[c] = events[k].aqcd + scale2[c] = events[k].scale ** 2 + # PY_GET_DENSITY_BATCH(PDGS, PROCID, P, POS, ALLOW_HEL, ALPHAS, + # SCALE2, NBATCH, NEXT) + inter = self.f2py_module.py_get_density_batch(pdgs=all_pdgs, + procid=-1, + p=P, + pos=position, + allow_hel=allow_hel, + alphas=alphas, + scale2=scale2, + nbatch=nbatch, + next=next_) + for c, k in enumerate(idx): + out[k] = madspin.DensityMatrix(inter[:, c], + n_changing, + allow_hel, + dimension) + return out + + def get_inter_value(self,event,nhel): """routine to return all the possible inter for an event""" @@ -2748,6 +3251,16 @@ def get_pdir(self,event): # the lookup with KeyError, e.g. ((-2, 2), (21, 23)) when the table # is keyed by ((-81, 81), (21, 23)). tag, order = event.get_tag_and_order(self._revert_merged or None) + # The answer is a property of the flavour tag alone, and a run sees only + # a handful of tags while calling this once per decaying particle per + # trial (123k times for 10k events). + event_tag = tag + try: + return self._pdir_cache[event_tag] + except AttributeError: + self._pdir_cache = {} + except KeyError: + pass try: orig_order = self.all_me[tag]['order'] except Exception: @@ -2762,6 +3275,9 @@ def get_pdir(self,event): pdir = self.all_me[tag]['pdir'] prefix, pos = self.pdg2prefix[tuple(list(orig_order[0]) + list(orig_order[1]))] #misc.sprint(f"get_pdir: pdir = {pdir} , orig_order = {orig_order} , prefix = {prefix}") + # Cache under the tag we were asked about, not the anti-particle tag the + # fallback above may have rewritten it to. + self._pdir_cache[event_tag] = (pdir, orig_order, prefix, pos) return pdir,orig_order, prefix, pos model_init = True diff --git a/madgraph/iolibs/export_cpp.py b/madgraph/iolibs/export_cpp.py index 354815b36..73d15dba1 100755 --- a/madgraph/iolibs/export_cpp.py +++ b/madgraph/iolibs/export_cpp.py @@ -3265,15 +3265,18 @@ def finalize(self, matrix_elements=None, history='', *args, **kwargs): with open(file_name, 'w') as f: json.dump(self.process_info, f) - # Generate Cards/run_card.toml from the template, filling in - # process-dependent defaults (mirrors the LO run_card.dat logic). - self.create_run_card(matrix_elements, history) - # SubProcesses/proc_characteristics: needed by the CommonRunCmd-based # post-processing driver (get_characteristics) so that the madevent # tool interface can run on this directory. + # NB: this must come *before* create_run_card: it is what fills in + # self.proc_characteristic, and the run_card defaults are derived from + # it (e.g. ninitial == 1 -> a decay, which gets no cuts at all). self.create_proc_characteristics(matrix_elements) + # Generate Cards/run_card.toml from the template, filling in + # process-dependent defaults (mirrors the LO run_card.dat logic). + self.create_run_card(matrix_elements, history) + # Cards/me5_configuration.txt: read by CommonRunCmd.set_configuration. # Point it at the MG5 install so tool paths (pythia8, etc.) and the # cluster/run-mode settings resolve from the central configuration. diff --git a/madgraph/iolibs/export_mg7.py b/madgraph/iolibs/export_mg7.py index a6cc4f051..2d07f564e 100644 --- a/madgraph/iolibs/export_mg7.py +++ b/madgraph/iolibs/export_mg7.py @@ -34,17 +34,33 @@ def generate_process_files(self): super().generate_process_files() def set_topology(self): + """Name every external leg i/o and record the initial/final pdgs. + + Two initial legs for a collision, one for a decay (``t > b w+, ...``, + which MadSpin hands over as a single flattened matrix element). Legs are + numbered 1..n with the initial state first, so the outgoing offset is + the number of initial legs. + """ self.edge_names = {} - self.incoming = [None] * 2 - self.outgoing = [None] * (len(self.legs) - 2) + self.n_initial = sum(1 for leg in self.legs if not leg.get("state")) + self.incoming = [None] * self.n_initial + self.outgoing = [None] * (len(self.legs) - self.n_initial) for leg in self.legs: number = leg.get("number") if leg.get("state"): - self.edge_names[number] = f"o{number - 3}" - self.outgoing[number - 3] = leg.get("id") + index = number - self.n_initial - 1 + self.edge_names[number] = f"o{index}" + self.outgoing[index] = leg.get("id") else: self.edge_names[number] = f"i{number - 1}" self.incoming[number - 1] = leg.get("id") + if any(pdg is None for pdg in self.incoming + self.outgoing): + raise AssertionError( + "external legs of %s are not numbered 1..%d with the initial " + "state first: %s" % ( + self.name, len(self.legs), + [(leg.get("number"), leg.get("state")) for leg in self.legs]) + ) def expand_flavors_over_processes(self): """Add the flavors that live in the *processes* mapped onto this matrix @@ -77,12 +93,15 @@ def expand_flavors_over_processes(self): self.all_flavors = [self.all_flavors[0] * len(pdg_lists)] def set_flavor_indices(self): + # Flavor combinations are grouped by their initial state: the launcher + # picks one initial state (PDF-weighted), then a final state within it. + # A decay has a single initial leg to group on, not a beam pair. self.all_flavors_same_initial = [] self.all_flavors_indices = [] for i, flavors in enumerate(self.all_flavors_pdgs): flv_dict = defaultdict(list) for flv in flavors: - flv_dict[(flv[0], flv[1])].append(flv) + flv_dict[tuple(flv[:self.n_initial])].append(flv) indices = [] for flv in flv_dict.values(): indices.append(len(self.all_flavors_same_initial)) @@ -245,15 +264,16 @@ def get_subprocess_info(self, proc_dir, lib_me_path): # "u q > u q" (q = u d) carry the same merged pdg (81), yet leg 1 is # fixed to u, so "d u > u d" is not part of the process and mirroring the # u d flavor would double count it. - same_initial_multiparticle = \ + # A decay has a single initial leg, so there is no beam swap to mirror. + same_initial_multiparticle = self.n_initial == 2 and \ self.matrix_element.get("processes")[0].has_same_initial_multiparticle() flavors = [ { "index": index, "options": options, - "mirror": has_mirror_all or ( + "mirror": self.n_initial == 2 and (has_mirror_all or ( same_initial_multiparticle and options[0][0] != options[0][1] - ) + )) } for index, options in self.all_flavors_same_initial ] diff --git a/madgraph/iolibs/template_files/f2py_wrapper_all.inc b/madgraph/iolibs/template_files/f2py_wrapper_all.inc index 527056ae8..7f9864399 100644 --- a/madgraph/iolibs/template_files/f2py_wrapper_all.inc +++ b/madgraph/iolibs/template_files/f2py_wrapper_all.inc @@ -45,7 +45,50 @@ CF2PY integer, intent(hide), depend(PDGS) :: NPDG = len(PDGS) DOUBLE COMPLEX INTER(N_COMB*(N_COMB+1)/2) call %(f2py_prefix)sf77_density(PDGS, NPDG, PROCID, P, POS, N_CHANGING, ALLOW_HEL, N_COMB, ALPHAS, SCALE2, INTER) - return + return + end + + SUBROUTINE %(f2py_prefix)sPY_GET_DENSITY_BATCH(PDGS, PROCID, P, POS, ALLOW_HEL, ALPHAS, SCALE2, NBATCH, NEXT, INTER, N_CHANGING, N_COMB, NPDG) +C PY_GET_DENSITY for NBATCH phase-space points in one call. +C +C The density itself is cheap; crossing into fortran is not. MadSpin's +C max-weight scan evaluates hundreds of decay configurations against the +C same production event, so paying the f2py entry once per point dominated +C the cost. Everything is per point except POS/ALLOW_HEL, which describe +C the helicity structure and are shared: PDGS is per point because a decay +C pool mixes flavours, and ALPHAS/SCALE2 because nothing guarantees the +C points share a scale. + IMPLICIT NONE +CF2PY double precision, intent(in), dimension(0:3,next,nbatch) :: p +CF2PY integer, intent(in), dimension(npdg,nbatch) :: pdgs +CF2PY integer, intent(in) :: procid +CF2PY integer, INTENT(IN) :: ALLOW_HEL(N_CHANGING*N_COMB) +CF2PY integer, INTENT(IN) :: POS(N_CHANGING) +CF2PY double precision INTENT(IN), dimension(nbatch) :: ALPHAS +CF2PY double precision INTENT(IN), dimension(nbatch) :: SCALE2 +CF2PY integer, intent(in) :: NBATCH +CF2PY integer, intent(in) :: NEXT +CF2PY double complex INTENT(OUT), dimension(N_COMB*(N_COMB+1)/2,NBATCH) :: INTER +CF2PY integer, intent(hide), depend(allow_hel, pos) :: N_COMB = len(ALLOW_HEL)/len(pos) +CF2PY integer, intent(hide), depend(pos) :: N_CHANGING = len(pos) +CF2PY integer, intent(hide), depend(PDGS) :: NPDG = len(PDGS) + + INTEGER NBATCH, NEXT, NPDG, PROCID, N_CHANGING, N_COMB + INTEGER PDGS(NPDG, NBATCH) + INTEGER POS(N_CHANGING) + DOUBLE PRECISION ALPHAS(NBATCH), SCALE2(NBATCH) + DOUBLE PRECISION P(0:3, NEXT, NBATCH) + INTEGER ALLOW_HEL(N_CHANGING*N_COMB) + DOUBLE COMPLEX INTER(N_COMB*(N_COMB+1)/2, NBATCH) + INTEGER IBATCH + + DO IBATCH = 1, NBATCH + call %(f2py_prefix)sf77_density(PDGS(1,IBATCH), NPDG, PROCID, + $ P(0,1,IBATCH), POS, N_CHANGING, ALLOW_HEL, N_COMB, + $ ALPHAS(IBATCH), SCALE2(IBATCH), INTER(1,IBATCH)) + ENDDO + + return end SUBROUTINE %(f2py_prefix)sPY_GET_ALL_INTER(PDGS, PROCID, NPDG, P, POS,N_CHANGING, ALLOW_HEL, N_COMB, INTER) diff --git a/madgraph/iolibs/template_files/mg7/madevent.py b/madgraph/iolibs/template_files/mg7/madevent.py index 9399c2ad5..36488c8c7 100644 --- a/madgraph/iolibs/template_files/mg7/madevent.py +++ b/madgraph/iolibs/template_files/mg7/madevent.py @@ -229,6 +229,77 @@ def load_cards(self) -> None: self.param_card = ParamCard(self.param_card_path) with open(os.path.join("SubProcesses", "subprocesses.json")) as f: self.subprocess_data = json.load(f) + self.init_decay_mode() + + def init_decay_mode(self) -> None: + """Decide whether this directory is a decay (1 -> n) or a collision. + + Read off the exported process rather than the run card: the two cannot + then disagree. MadSpin generates its decay matrix elements this way. + """ + incoming_counts = { + len(clean_pids(meta["incoming"])) for meta in self.subprocess_data + } + if incoming_counts - {1, 2}: + raise ValueError( + f"processes with {sorted(incoming_counts)} incoming particles " + "are not supported" + ) + if len(incoming_counts) > 1: + raise ValueError( + "cannot mix decays and collisions in one output directory" + ) + self.is_decay = incoming_counts == {1} + if not self.is_decay: + self.decaying_mass = None + return + + masses = { + self.get_mass(clean_pids(meta["incoming"])[0]) + for meta in self.subprocess_data + } + if len(masses) > 1: + raise ValueError( + f"decaying particles have different masses: {sorted(masses)}" + ) + self.decaying_mass = masses.pop() + if self.decaying_mass <= 0.0: + raise ValueError("the decaying particle must have a non-zero mass") + self.drop_closed_channels() + + def drop_closed_channels(self) -> None: + """Remove subprocesses the decaying particle is too light to produce. + + A multiparticle decay definition enumerates every vertex the model + allows, closed ones included: "t > b w+, w+ > all all" yields + t > b t b~ (and b W+ Z, b W+ h, ...), which need more mass than the top + has. Their partial width is exactly zero, but the phase-space mapping + has no physical point to hand back -- the invariant's lower bound ends + up above its upper bound -- so it produces NaN momenta and poisons the + whole integral. Drop them here instead. + """ + kept, dropped = [], [] + for meta in self.subprocess_data: + total = sum( + self.get_mass(pid) for pid in clean_pids(meta["outgoing"]) + ) + if total < self.decaying_mass: + kept.append(meta) + else: + dropped.append((meta["outgoing"], total)) + if dropped: + for outgoing, total in dropped: + logger.info( + "skipping closed decay channel -> %s (needs %.4g GeV, " + "the decaying particle has %.4g GeV)", + outgoing, total, self.decaying_mass, + ) + if not kept: + raise ValueError( + "every decay channel is kinematically closed: the decaying " + f"particle's mass is {self.decaying_mass} GeV" + ) + self.subprocess_data = kept def init_backend(self) -> None: ms.set_simd_vector_size(self.run_card["run"]["simd_vector_size"]) @@ -384,8 +455,14 @@ def ensure_pdf_set(self, pdf_set: str) -> None: def init_beam(self) -> None: beam_args = self.run_card["beam"] - self.e_cm = beam_args["e_cm"] - self.leptonic = beam_args["leptonic"] + if self.is_decay: + # No beams: the total energy is the decaying particle's mass, and + # "leptonic" is what the mappings call "no parton luminosity". + self.e_cm = self.decaying_mass + self.leptonic = True + else: + self.e_cm = beam_args["e_cm"] + self.leptonic = beam_args["leptonic"] dynamical_scales = { "transverse_energy": ms.EnergyScale.transverse_energy, @@ -405,6 +482,22 @@ def init_beam(self) -> None: fact_scale1=beam_args["fact_scale1"], fact_scale2=beam_args["fact_scale2"], ) + if self.is_decay: + # One scale is available for a decay -- the decaying mass -- so use + # it, fixed, whatever the card asks for. + self.scale_kwargs.update( + ren_scale_fixed=True, + fact_scale_fixed=True, + ren_scale=self.decaying_mass, + fact_scale1=self.decaying_mass, + fact_scale2=self.decaying_mass, + ) + self.pdf_grid = None + self.alphas_grid = ms.AlphaSGrid(self.write_fixed_alphas_info()) + for context in self.contexts: + self.alphas_grid.initialize_globals(context) + self.running_coupling = ms.RunningCoupling(self.alphas_grid) + return pdf_set = beam_args["pdf"] self.ensure_pdf_set(pdf_set) @@ -417,6 +510,35 @@ def init_beam(self) -> None: self.alphas_grid.initialize_globals(context) self.running_coupling = ms.RunningCoupling(self.alphas_grid) + def write_fixed_alphas_info(self) -> str: + """Write a minimal LHAPDF ``.info`` holding a constant alpha_s. + + A decay has no beams, so there is no PDF set to take alpha_s from -- + and demanding one (possibly downloading it) just to evaluate a coupling + would be absurd. The renormalisation scale of a decay is fixed at the + decaying mass anyway, so a constant alpha_s is the right answer, not an + approximation: take it from the param card, exactly as the Fortran + decay matrix elements do. + """ + # SMINPUTS entry 3 is alpha_s(m_Z), the same value the Fortran + # parameter setup feeds to G. + alpha_s = float(self.param_card.get_value("sminputs", 3)) + path = os.path.join(self.run_path, "fixed_alphas.info") + # AlphaSGrid interpolates in log(q^2) across the grid, so give it a + # comfortable number of nodes rather than the bare minimum of three. + q_values = [10 ** (i / 4.0) for i in range(-4, 21)] + with open(path, "w") as f: + f.write("SetDesc: constant alpha_s for a decay (no beams)\n") + f.write("AlphaS_Qs: [%s]\n" % ", ".join(f"{q:g}" for q in q_values)) + f.write( + "AlphaS_Vals: [%s]\n" + % ", ".join(f"{alpha_s:g}" for _ in q_values) + ) + logger.info( + "decay mode: fixed alpha_s = %g at mu = %g GeV", alpha_s, self.e_cm + ) + return path + def init_generator_config(self) -> None: run_args = self.run_card["run"] gen_args = self.run_card["generation"] @@ -569,6 +691,15 @@ def survey_phasespaces( def survey(self) -> None: phasespace_mode = self.run_card["phasespace"]["mode"] + if self.is_decay and phasespace_mode != "multichannel": + # The flat mapping is built from a synthetic two-incoming diagram, + # which a decay has no counterpart for. Decays have few channels, so + # the multichannel phase space is the right one regardless. + logger.info( + "decay mode: using the multichannel phase space instead of " + "'%s'", phasespace_mode + ) + phasespace_mode = "multichannel" if phasespace_mode == "multichannel": self.phasespaces = [ subproc.build_multichannel_phasespace() @@ -878,6 +1009,12 @@ def _beam_info(self): beam particle, so hadronic beams are protons (2212); leptonic beams are the incoming leptons themselves.""" half_e = float(self.e_cm) / 2. + if self.is_decay: + # No beams. LHE has no way to say that, so report the decaying + # particle at rest as a single "beam"; the second slot is empty. + data = self.subprocess_data[0] + pdg = clean_pids(data["incoming"])[0] + return [pdg, 0], [float(self.e_cm), 0.0] if not self.leptonic: # hadronic collider: proton beams (p-pbar is not distinguished) return [2212, 2212], [half_e, half_e] @@ -1062,14 +1199,29 @@ def get_width(self, pid: int) -> float: return self.param_card.get_value("width", pid) +# Flavor-merged legs carry a group id instead of a pdg. Every member of a group +# shares the same mass -- that is what makes them mergeable -- so any member is +# a valid representative for the kinematics. +_MERGED_PID_REPRESENTATIVE = { + 81: 1, # light quarks d u s c + 82: 11, # charged leptons e mu + 83: 12, # neutrinos ve vm vt +} + + def clean_pids(pids: list[int]) -> list[int]: pids_out = [] for pid in pids: pid = abs(pid) - if pid == 81: - pid = 1 - if pid == 82: - pid = 11 + if pid in _MERGED_PID_REPRESENTATIVE: + pid = _MERGED_PID_REPRESENTATIVE[pid] + elif 81 <= pid <= 99: + # Reserved for flavor merging. Anything outside this window is a + # real pdg (BSM models use codes in the millions), so let it pass. + raise ValueError( + f"unknown flavor-merged particle id {pid}; add its " + "representative to _MERGED_PID_REPRESENTATIVE" + ) pids_out.append(pid) return pids_out @@ -1587,6 +1739,9 @@ def build_integrands( pdf_arg = None if self.process.leptonic else ms.CachedPdf() cross_section = ms.DifferentialCrossSection( matrix_element=matrix_element, + # For a decay this is the decaying particle's mass, and the flux + # becomes 1/(2M): the result is a partial width in GeV, not a cross + # section in pb. cm_energy=self.process.e_cm, running_coupling=None, energy_scale=ms.CachedScale(), @@ -1594,6 +1749,7 @@ def build_integrands( pdf1=pdf_arg, pdf2=pdf_arg, input_momentum_fraction=True, + decay=self.process.is_decay, ) partial_weights = self.process.run_card["generation"]["systematics"] integrands = [] @@ -2184,7 +2340,12 @@ def run_lhe_postprocessing(process) -> None: return log = logging.getLogger('madevent') - if cfg.get('systematics'): + if cfg.get('systematics') and getattr(process, 'is_decay', False): + # Scale and PDF variations are a beam quantity. A decay has neither, so + # systematics can only fail here ("not supported for pdlabel=none") -- + # and it is not free: MadSpin reruns the launcher for every pool refill. + log.info("decay mode: skipping systematics (no beams to vary)") + elif cfg.get('systematics'): try: _run_systematics(lhe_path, cfg, log) except Exception as error: diff --git a/madgraph/various/banner.py b/madgraph/various/banner.py index 5d5b0c63d..16330df76 100755 --- a/madgraph/various/banner.py +++ b/madgraph/various/banner.py @@ -7114,9 +7114,13 @@ def create_default_for_process(self, proc_characteristic, history, proc_def): # e.g. photon/neutrino initiated: also no proton PDF self['beam']['leptonic'] = True - # 1 -> N decay: no cuts at all + # 1 -> N decay: no cuts at all. A partial width is an inclusive + # quantity, so any kinematic cut biases it low (the collider defaults + # cost ~3% on t > b w+, w+ > e+ ve). The Breit-Wigner cutoff in + # [phasespace] is deliberately left alone: it is a sampling range for + # the off-shell propagators, not a cut on the final state. if proc_characteristic and proc_characteristic['ninitial'] == 1: - self.dynamic_sections['cuts'] = collections.OrderedDict() + self.remove_all_cut() # site/user defaults win (this has to be LAST, like the LO run_card) if self.default_run_card and os.path.exists(self.default_run_card): diff --git a/madgraph/various/lhe_parser.py b/madgraph/various/lhe_parser.py index 6f93ddc66..9dfb16ab2 100755 --- a/madgraph/various/lhe_parser.py +++ b/madgraph/various/lhe_parser.py @@ -48,7 +48,10 @@ logger = logging.getLogger("madgraph.lhe_parser") -_ENABLE_LHE_TIMERS = False +# Parsing/writing timers. Off by default (the perf_counter calls are not free +# on the per-particle path); set MG_LHE_TIMERS=1 in the environment to collect +# them. MadSpin prints the summary through _log_lhe_timers at the end of a run. +_ENABLE_LHE_TIMERS = os.environ.get('MG_LHE_TIMERS', '') not in ('', '0', 'false', 'False') _LHE_TIMERS = collections.defaultdict(float) _LHE_TIMER_COUNTS = collections.defaultdict(int) @@ -2773,22 +2776,48 @@ def boost(self, filter=None): if filter is None: filter = lambda p: p.status==-1 - if not isinstance(filter, FourMomentum): + # Inline of FourMomentum.boost for the whole event. MadSpin calls this + # once per decay per accept/reject trial, and the original built two + # FourMomentum objects per particle only to copy four floats back out. + # The boost vector is shared by every particle, so its norm, mass and + # the two divisions by them are loop invariant; hoisting them leaves + # per particle only multiplies and adds. + # + # The helas sign flip is applied to locals rather than to a copy of the + # boost momentum, which saves the copy and leaves the caller's object + # untouched exactly as the copy did. + if isinstance(filter, FourMomentum): + bpx, bpy, bpz, bE = -filter.px, -filter.py, -filter.pz, filter.E + else: pboost = FourMomentum() for p in self: if list(filter(p)): pboost += p - else: - pboost = FourMomentum(filter) + bpx, bpy, bpz, bE = -pboost.px, -pboost.py, -pboost.pz, pboost.E + + pnorm = bpx * bpx + bpy * bpy + bpz * bpz + if pnorm and len(self): + # Rounds differently from FourMomentum.boost, which divides per + # particle and forms the mass as E^2-px^2-py^2-pz^2 rather than + # E^2-pnorm. Measured over 40k components of real decay events the + # worst relative difference is 1.2e-12, in the cases where E+k*s3 + # cancels; typical components agree to ~1e-16. That is far below + # any physical sensitivity here. + mass = math.sqrt(max(bE * bE - pnorm, 0.)) + inv_mass = 1.0 / mass + k = (bE - mass) / pnorm + for p in self: + px, py, pz, E = p.px, p.py, p.pz, p.E + s3product = px * bpx + py * bpy + pz * bpz + lf = (E + k * s3product) * inv_mass + p.E = (E * bE + s3product) * inv_mass + p.px = px + bpx * lf + p.py = py + bpy * lf + p.pz = pz + bpz * lf + elif not pnorm: + for p in self: + p.E, p.px, p.py, p.pz = bE, bpx, bpy, bpz - # change sign of three-component due to helas convention - pboost.px *=-1 - pboost.py *=-1 - pboost.pz *=-1 - for p in self: - b= FourMomentum(p).boost(pboost) - p.E, p.px, p.py, p.pz = b.E, b.px, b.py, b.pz - return self def check(self): diff --git a/madgraph/various/misc.py b/madgraph/various/misc.py index 15e28d903..d83829819 100755 --- a/madgraph/various/misc.py +++ b/madgraph/various/misc.py @@ -1342,7 +1342,17 @@ def configure_gzip(configuration=None): if configuration['nb_core'] is not None: _gzip_tool_max_cores = configuration['nb_core'] -def gzip(path, stdout=None, error=True, forceexternal=False): +# Compression level for the in-process branch of gzip() below. The gzip module +# defaults to 9, which is a poor trade here: on a 172 MB LHE, level 9 takes +# 18.6 s against 4.5 s at level 6, and buys 4% (38.1 MB against 39.7 MB). Level +# 6 is also what the external tool this function shells out to for large files +# uses, so the two branches now agree instead of compressing the same data +# differently depending on its size. +GZIP_COMPRESSLEVEL = 6 + + +def gzip(path, stdout=None, error=True, forceexternal=False, + compresslevel=GZIP_COMPRESSLEVEL): """ a standard replacement for os.system('gzip %s ' % path)""" # For large files (>256M), it is faster and safer to use a separate tool. @@ -1364,8 +1374,11 @@ def gzip(path, stdout=None, error=True, forceexternal=False): stdout = "%s.gz" % stdout try: - with ziplib.open(stdout, 'wb') as f: - f.write(open(path).read().encode()) + # Stream it: reading the whole file in as a str and encoding it made a + # 172 MB LHE cost two extra full-size copies in memory. + with open(path, 'rb') as fsock, \ + ziplib.open(stdout, 'wb', compresslevel=compresslevel) as f: + shutil.copyfileobj(fsock, f, 4 * 1024 * 1024) except OverflowError: gzip(path, stdout, error=error, forceexternal=True) except Exception: diff --git a/madspace/include/madspace/driver/lhe_output.hpp b/madspace/include/madspace/driver/lhe_output.hpp index 630175a22..7e9bc9ac9 100644 --- a/madspace/include/madspace/driver/lhe_output.hpp +++ b/madspace/include/madspace/driver/lhe_output.hpp @@ -98,6 +98,9 @@ class LHECompleter { std::size_t color_offset, pdg_id_offset, helicity_offset, mass_offset; std::size_t particle_count, color_count, flavor_count; std::size_t diagram_count, helicity_count; + // 2 for a collision, 1 for a decay. Decides which leading particles are + // written as initial state and what the outgoing ones point at. + std::size_t incoming_count; }; struct PropagatorData { int pdg_id; diff --git a/madspace/include/madspace/phasespace/cross_section.hpp b/madspace/include/madspace/phasespace/cross_section.hpp index 19339e409..c43081917 100644 --- a/madspace/include/madspace/phasespace/cross_section.hpp +++ b/madspace/include/madspace/phasespace/cross_section.hpp @@ -20,7 +20,12 @@ class DifferentialCrossSection : public FunctionGenerator { const nested_vector2& pid_options = {}, const std::variant& pdf1 = std::monostate{}, const std::variant& pdf2 = std::monostate{}, - bool input_momentum_fraction = true + bool input_momentum_fraction = true, + // Decay of a single incoming particle instead of a collision: the + // differential rate is |M|^2 / (2 M) with M = cm_energy, in GeV, rather + // than a hadronic cross section in pb. There are no beams, hence no + // momentum fractions and no PDFs. + bool decay = false ); const nested_vector2& pid_options() const { return _pid_options; } @@ -43,6 +48,7 @@ class DifferentialCrossSection : public FunctionGenerator { double _e_cm; std::variant _energy_scale; bool _input_momentum_fraction; + bool _decay; }; } // namespace madspace diff --git a/madspace/include/madspace/phasespace/phasespace.hpp b/madspace/include/madspace/phasespace/phasespace.hpp index b1827565d..285162760 100644 --- a/madspace/include/madspace/phasespace/phasespace.hpp +++ b/madspace/include/madspace/phasespace/phasespace.hpp @@ -38,12 +38,17 @@ class PhaseSpaceMapping : public Mapping { const std::optional>& color_order = std::nullopt ); - std::size_t random_dim() const { - return 3 * _topology.outgoing_masses().size() - (_leptonic ? 4 : 2); + // A 1 -> n decay and a leptonic (fixed-s) 2 -> n collision both have 3n-4 + // degrees of freedom; a hadronic 2 -> n adds the beam momentum fractions, + // i.e. one further sampled invariant on top of s_hat -> 3n-2. + static std::size_t random_dim_for(const Topology& topology, bool leptonic) { + return 3 * topology.outgoing_masses().size() - + ((leptonic || topology.is_decay()) ? 4 : 2); } + std::size_t random_dim() const { return random_dim_for(_topology, _leptonic); } std::size_t discrete_dim() const override { return _n_discrete; } std::size_t particle_count() const { - return _topology.outgoing_masses().size() + 2; + return _topology.outgoing_masses().size() + _topology.incoming_masses().size(); } std::size_t channel_count() const { return _permutations.size(); } diff --git a/madspace/include/madspace/phasespace/topology.hpp b/madspace/include/madspace/phasespace/topology.hpp index c023e970d..bb54e0f8c 100644 --- a/madspace/include/madspace/phasespace/topology.hpp +++ b/madspace/include/madspace/phasespace/topology.hpp @@ -43,8 +43,10 @@ class Diagram { const std::vector& outgoing_masses() const { return _outgoing_masses; } const std::vector& propagators() const { return _propagators; } const std::vector& vertices() const { return _vertices; } - const std::array& incoming_vertices() const { return _incoming_vertices; }; + // One entry per incoming particle: two for a collision, one for a decay. + const std::vector& incoming_vertices() const { return _incoming_vertices; }; const std::vector& outgoing_vertices() const { return _outgoing_vertices; }; + bool is_decay() const { return _incoming_masses.size() == 1; } const std::vector>& propagator_vertices() const { return _propagator_vertices; } @@ -54,7 +56,7 @@ class Diagram { std::vector _outgoing_masses; std::vector _propagators; std::vector _vertices; - std::array _incoming_vertices; + std::vector _incoming_vertices; std::vector _outgoing_vertices; std::vector> _propagator_vertices; }; @@ -98,6 +100,9 @@ class Topology { } const std::vector& incoming_masses() const { return _incoming_masses; } const std::vector& outgoing_masses() const { return _outgoing_masses; } + // A decay (single incoming particle) has no t-channel chain and a root + // virtuality fixed by the decaying particle's mass. + bool is_decay() const { return _incoming_masses.size() == 1; } std::vector, double, double>> propagator_momentum_terms(bool only_decays = false) const; std::string to_string() const; diff --git a/madspace/src/driver/lhe_output.cpp b/madspace/src/driver/lhe_output.cpp index 3f96379b7..f010f7b5b 100644 --- a/madspace/src/driver/lhe_output.cpp +++ b/madspace/src/driver/lhe_output.cpp @@ -200,10 +200,16 @@ void LHECompleter::init_propagator_data( resonant_prop_indices.clear(); resonant_prop_indices.resize(decay_count, -1); + // The permutation runs over all external legs, incoming first, so the + // outgoing ones start at the incoming count -- 2 for a collision, 1 for a + // decay. Getting this wrong silently reads another particle's color flow. for (auto [index, mass, perm_index] : zip(topo.outgoing_indices(), topo.outgoing_masses(), - std::span(permutation.begin() + 2, permutation.end()))) { + std::span( + permutation.begin() + topo.incoming_masses().size(), + permutation.end() + ))) { e_min.at(index) = mass; momentum_masks.at(index) = 1 << perm_index; for (std::size_t i = 0; std::size_t color_index : colors) { @@ -377,6 +383,7 @@ LHECompleter::LHECompleter( .flavor_count = args.pdg_ids.size(), .diagram_count = diagram_count, .helicity_count = args.helicities.size(), + .incoming_count = args.topologies.at(0).incoming_masses().size(), }); helicity_offset += particle_count * args.helicities.size(); @@ -415,6 +422,13 @@ void LHECompleter::complete_event_data( event.process_id = subproc_data.process_id; + // Number of leading entries in the event record that are initial state: + // 2 for a collision, 1 for a decay. Everything that indexes past the + // initial state -- where resonances get inserted, what the outgoing + // particles' mothers are, how far the momentum masks are shifted -- is + // offset by this rather than by a hard-coded 2. + const std::size_t n_in = subproc_data.incoming_count; + std::size_t color_offset = subproc_data.color_offset + subproc_data.particle_count * color_index; std::size_t helicity_offset = @@ -431,14 +445,16 @@ void LHECompleter::complete_event_data( std::tie(particle.color, particle.anti_color) = _colors.at(color_offset + particle_index); particle.pdg_id = _pdg_ids.at(pdg_offset + particle_index); - if (particle_index < 2) { + if (particle_index < n_in) { particle.status_code = -1; particle.mother1 = 0; particle.mother2 = 0; } else { particle.status_code = 1; particle.mother1 = 1; - particle.mother2 = 2; + // A decay has a single mother, which LHE spells as + // mother1 == mother2 rather than a (1, 2) range. + particle.mother2 = static_cast(n_in); } particle.mass = _masses.at(mass_offset + particle_index); particle.lifetime = 0; @@ -486,7 +502,7 @@ void LHECompleter::complete_event_data( .pdg_id = propagator.pdg_id, .status_code = 2, .mother1 = 1, - .mother2 = 2, + .mother2 = static_cast(n_in), .color = color, .anti_color = anti_color, .px = px, @@ -501,7 +517,7 @@ void LHECompleter::complete_event_data( ++prop_index; } event.particles.insert( - event.particles.begin() + 2, new_particles.rbegin(), new_particles.rend() + event.particles.begin() + n_in, new_particles.rbegin(), new_particles.rend() ); for (std::size_t prop_index = prop_count, res_index = 0; auto& propagator : std::views::reverse( @@ -517,21 +533,21 @@ void LHECompleter::complete_event_data( child_prop_index >= 0; --child_prop_index) { if (child_prop_mask & (1 << child_prop_index)) { - auto& child_particle = event.particles.at(child_res_index + 2); - child_particle.mother1 = res_index + 3; - child_particle.mother2 = res_index + 3; + auto& child_particle = event.particles.at(child_res_index + n_in); + child_particle.mother1 = res_index + n_in + 1; + child_particle.mother2 = res_index + n_in + 1; ++child_res_index; } } - int momentum_mask = propagator.momentum_mask >> 2; + int momentum_mask = propagator.momentum_mask >> n_in; for (auto& particle : std::span( - event.particles.begin() + 2 + new_particles.size(), + event.particles.begin() + n_in + new_particles.size(), event.particles.end() )) { if (momentum_mask & 1) { - particle.mother1 = res_index + 3; - particle.mother2 = res_index + 3; + particle.mother1 = res_index + n_in + 1; + particle.mother2 = res_index + n_in + 1; } momentum_mask >>= 1; } @@ -613,6 +629,7 @@ void madspace::to_json( subproc_data.flavor_count, subproc_data.diagram_count, subproc_data.helicity_count, + subproc_data.incoming_count, }; } @@ -630,6 +647,9 @@ void madspace::from_json( .flavor_count = j.at(7).get(), .diagram_count = j.at(8).get(), .helicity_count = j.at(9).get(), + // Gridpacks written before decays were supported carry no entry here + // and are always collisions. + .incoming_count = j.size() > 10 ? j.at(10).get() : 2, }; } diff --git a/madspace/src/phasespace/cross_section.cpp b/madspace/src/phasespace/cross_section.cpp index 9ed44d972..f5f46b31c 100644 --- a/madspace/src/phasespace/cross_section.cpp +++ b/madspace/src/phasespace/cross_section.cpp @@ -10,7 +10,8 @@ DifferentialCrossSection::DifferentialCrossSection( const nested_vector2& pid_options, const std::variant& pdf1, const std::variant& pdf2, - bool input_momentum_fraction + bool input_momentum_fraction, + bool decay ) : FunctionGenerator( "DifferentialCrossSection", @@ -60,7 +61,18 @@ DifferentialCrossSection::DifferentialCrossSection( _running_coupling(running_coupling), _e_cm(cm_energy), _energy_scale(energy_scale), - _input_momentum_fraction(input_momentum_fraction) { + _input_momentum_fraction(input_momentum_fraction), + _decay(decay) { + if (decay) { + if (_has_pdf.at(0) || _has_pdf.at(1)) { + throw std::invalid_argument("a decay has no beams, so no PDFs"); + } + if (cm_energy <= 0.) { + throw std::invalid_argument( + "a decay needs the decaying particle's mass as cm_energy" + ); + } + } auto init_pdf = [&](auto& pdf, int index) { if (std::holds_alternative(pdf)) { std::vector pids; @@ -93,8 +105,15 @@ NamedVector DifferentialCrossSection::build_function_impl( } } - std::array x1x2; - if (_input_momentum_fraction) { + std::array x1x2{1., 1.}; + if (_decay) { + // No beams: nothing to extract momentum fractions from. The declared + // argument list still carries x1/x2 when input_momentum_fraction is on, + // so consume them and ignore their (unit) values. + if (_input_momentum_fraction) { + arg_index += 2; + } + } else if (_input_momentum_fraction) { x1x2 = {args.at(arg_index), args.at(arg_index + 1)}; arg_index += 2; } else { @@ -138,13 +157,23 @@ NamedVector DifferentialCrossSection::build_function_impl( throw std::runtime_error("matrix element missing in return values"); } std::size_t me_index = search - _matrix_element.outputs().begin(); - me_result.at(me_index) = fb.diff_cross_section( - x1x2.at(0), - x1x2.at(1), - pdf_outputs.at(0), - pdf_outputs.at(1), - me_result.at(me_index), - _e_cm * _e_cm - ); + if (_decay) { + // Differential decay rate: dGamma = |M|^2 dPhi_n / (2 M), in GeV. The + // flux is a compile-time constant here (M is fixed by the decaying + // particle), so this is a single multiply and needs no dedicated + // instruction. Deliberately no GeV^-2 -> pb conversion: a width is not + // a cross section. + me_result.at(me_index) = + fb.mul(me_result.at(me_index), 1. / (2. * _e_cm)); + } else { + me_result.at(me_index) = fb.diff_cross_section( + x1x2.at(0), + x1x2.at(1), + pdf_outputs.at(0), + pdf_outputs.at(1), + me_result.at(me_index), + _e_cm * _e_cm + ); + } return me_result; } diff --git a/madspace/src/phasespace/phasespace.cpp b/madspace/src/phasespace/phasespace.cpp index f95f2e165..7de53dba1 100644 --- a/madspace/src/phasespace/phasespace.cpp +++ b/madspace/src/phasespace/phasespace.cpp @@ -134,7 +134,7 @@ PhaseSpaceMapping::PhaseSpaceMapping( NamedVector in{ {"random", batch_float_array( - 3 * topology.outgoing_masses().size() - (leptonic ? 4 : 2) + PhaseSpaceMapping::random_dim_for(topology, leptonic) )} }; // Opt-in discrete channel: only declared when the t-channel strategy @@ -148,7 +148,10 @@ PhaseSpaceMapping::PhaseSpaceMapping( } return in; }(), - {{"momenta", batch_four_vec_array(topology.outgoing_masses().size() + 2)}, + {{"momenta", + batch_four_vec_array( + topology.outgoing_masses().size() + topology.incoming_masses().size() + )}, {"x1", batch_float}, {"x2", batch_float}}, permutations.size() > 1 @@ -156,14 +159,19 @@ PhaseSpaceMapping::PhaseSpaceMapping( : NamedVector{} ), _topology(topology), - _cuts(cuts.value_or(Cuts(topology.outgoing_masses().size() + 2))), + _cuts(cuts.value_or(Cuts( + topology.outgoing_masses().size() + topology.incoming_masses().size() + ))), _pi_factors( std::pow(2 * PI, 4 - 3 * static_cast(topology.outgoing_masses().size())) ), _sqrt_s_lab(cm_energy), _leptonic(leptonic), + // A decay has no beams to sample momentum fractions for: the root + // virtuality is fixed at the decaying particle's mass (passed as + // cm_energy) and there is no boost into a lab frame. _map_luminosity( - !leptonic && + !leptonic && !topology.is_decay() && (_topology.t_propagator_count() == 0 || t_channel_mode != PhaseSpaceMapping::chili) ), @@ -479,7 +487,15 @@ Mapping::Result PhaseSpaceMapping::build_forward_impl( }, [&](std::monostate) { auto [p1, p2] = fb.com_p_in(sqrt_s_hat); - p_ext = {p1, p2}; + if (_topology.is_decay()) { + // Single incoming particle, at rest in the frame the decay + // products are generated in: p_in = (M, 0, 0, 0), which is + // exactly the sum of the two back-to-back beam momenta + // com_p_in builds for sqrt(s_hat) = M. + p_ext = {fb.add(p1, p2)}; + } else { + p_ext = {p1, p2}; + } } }, _t_mapping @@ -570,7 +586,9 @@ Mapping::Result PhaseSpaceMapping::build_inverse_impl( for (auto [decay_index, mass, momentum] : zip(_topology.outgoing_indices(), _topology.outgoing_masses(), - std::span(p_ext.begin() + 2, p_ext.end()))) { + std::span( + p_ext.begin() + _topology.incoming_masses().size(), p_ext.end() + ))) { auto& data = decay_data.at(decay_index); data.mass = mass; data.mass2 = mass * mass; diff --git a/madspace/src/phasespace/topology.cpp b/madspace/src/phasespace/topology.cpp index 90fbac951..f41471245 100644 --- a/madspace/src/phasespace/topology.cpp +++ b/madspace/src/phasespace/topology.cpp @@ -230,11 +230,13 @@ Diagram::Diagram( _outgoing_masses(outgoing_masses), _propagators(propagators), _vertices(vertices), - _incoming_vertices{-1, -1}, + _incoming_vertices(incoming_masses.size(), -1), _outgoing_vertices(outgoing_masses.size(), -1), _propagator_vertices(propagators.size()) { - if (incoming_masses.size() != 2) { - throw std::invalid_argument("Diagram must have two incoming particles"); + if (incoming_masses.size() != 1 && incoming_masses.size() != 2) { + throw std::invalid_argument( + "Diagram must have one incoming particle (a decay) or two (a collision)" + ); } if (outgoing_masses.size() < 2) { throw std::invalid_argument( @@ -291,38 +293,56 @@ std::vector Topology::topologies(const Diagram& diagram) { std::vector t_vertices; std::vector lines_after_t; std::vector integration_order; - find_t_vertices( - diagram, - visited, - t_vertices, - lines_after_t, - integration_order, - topo._t_propagator_masses, - topo._t_propagator_widths, - diagram.incoming_vertices().at(1), - -1 - ); - - // sort by integration order and propagator mass, while preventing - // impossible integration orders - bool choose_low = false; - std::size_t index_low = 0, index_high = integration_order.size() - 1; - while (index_low != index_high) { - int order_low = integration_order.at(index_low); - int order_high = integration_order.at(index_high - 1); - double mass_low = topo._t_propagator_masses.at(index_low); - double mass_high = topo._t_propagator_masses.at(index_high - 1); - if (order_low != order_high) { - choose_low = order_low < order_high; - } else if (mass_low != mass_high) { // TODO: maybe smarter heuristic here? - choose_low = mass_low < mass_high; + if (diagram.incoming_masses().size() == 1) { + // Decay: there is no t-channel chain to find. find_t_vertices walks in + // from the *second* incoming particle and marks the vertices between + // the two beams; with a single incoming particle the vertex it attaches + // to is simply the root of a pure s-channel cascade, and every other + // line at that vertex is a child of the root decay. Hand those to the + // shared decay builder below and leave _t_integration_order empty, so + // t_propagator_count() == 0 and consumers take their no-t-channel path. + std::size_t root_vertex = diagram.incoming_vertices().at(0); + for (auto& line_ref : diagram.vertices().at(root_vertex)) { + if (line_ref.type() == Diagram::incoming) { + continue; + } + t_vertices.push_back(root_vertex); + lines_after_t.push_back(line_ref); } - if (choose_low) { - topo._t_integration_order.push_back(index_low); - ++index_low; - } else { - topo._t_integration_order.push_back(index_high - 1); - --index_high; + } else { + find_t_vertices( + diagram, + visited, + t_vertices, + lines_after_t, + integration_order, + topo._t_propagator_masses, + topo._t_propagator_widths, + diagram.incoming_vertices().at(1), + -1 + ); + + // sort by integration order and propagator mass, while preventing + // impossible integration orders + bool choose_low = false; + std::size_t index_low = 0, index_high = integration_order.size() - 1; + while (index_low != index_high) { + int order_low = integration_order.at(index_low); + int order_high = integration_order.at(index_high - 1); + double mass_low = topo._t_propagator_masses.at(index_low); + double mass_high = topo._t_propagator_masses.at(index_high - 1); + if (order_low != order_high) { + choose_low = order_low < order_high; + } else if (mass_low != mass_high) { // TODO: maybe smarter heuristic here? + choose_low = mass_low < mass_high; + } + if (choose_low) { + topo._t_integration_order.push_back(index_low); + ++index_low; + } else { + topo._t_integration_order.push_back(index_high - 1); + --index_high; + } } } @@ -450,8 +470,11 @@ std::vector, double, double>> Topology::propagator_momentum_terms(bool only_decays) const { std::vector, double, double>> ret; std::vector> decay_indices(_decays.size()); - std::size_t n_ext = _outgoing_masses.size() + 2; - std::size_t ext_index = 2; + // The external momenta are laid out incoming-first, so the outgoing ones + // start at n_in (2 for a collision, 1 for a decay). + std::size_t n_in = _incoming_masses.size(); + std::size_t n_ext = _outgoing_masses.size() + n_in; + std::size_t ext_index = n_in; for (std::size_t index : _outgoing_indices) { decay_indices.at(index).push_back(ext_index); ++ext_index; @@ -460,8 +483,9 @@ Topology::propagator_momentum_terms(bool only_decays) const { if (decay.index == 0) { if (_t_integration_order.size() == 0) { std::vector factors(n_ext); - factors.at(0) = 1; - factors.at(1) = 1; + for (std::size_t i = 0; i < n_in; ++i) { + factors.at(i) = 1; + } ret.push_back({factors, decay.mass, decay.width}); } } else if (decay.child_indices.size() != 0) { diff --git a/madspace/src/python/madspace.cpp b/madspace/src/python/madspace.cpp index 72c029e7c..2b8e3aa6b 100644 --- a/madspace/src/python/madspace.cpp +++ b/madspace/src/python/madspace.cpp @@ -1196,6 +1196,7 @@ PYBIND11_MODULE(_madspace_py, m) { std::monostate, PdfGrid, DifferentialCrossSection::CachedPdf>&, + bool, bool>(), py::arg("matrix_element"), py::arg("cm_energy"), @@ -1204,7 +1205,8 @@ PYBIND11_MODULE(_madspace_py, m) { py::arg("pid_options") = nested_vector2{}, py::arg("pdf1") = std::monostate{}, py::arg("pdf2") = std::monostate{}, - py::arg("input_momentum_fraction") = true + py::arg("input_momentum_fraction") = true, + py::arg("decay") = false ) .def("pid_options", &DifferentialCrossSection::pid_options) .def("matrix_element", &DifferentialCrossSection::matrix_element); diff --git a/madspace/tests/test_decay_topology.py b/madspace/tests/test_decay_topology.py new file mode 100644 index 000000000..a58fc9c9b --- /dev/null +++ b/madspace/tests/test_decay_topology.py @@ -0,0 +1,206 @@ +"""Phase-space mappings for a decay, i.e. a single incoming particle. + +A 1 -> n topology has no t-channel chain and a root virtuality fixed by the +decaying particle's mass, so it reuses the s-channel cascade the 2 -> n mappings +already build. These tests pin that down against results that are known +analytically: + + * the total phase-space volume, + * the external masses and momentum conservation, + * the decaying particle at rest, p_in = (M, 0, 0, 0), + * an exact forward/inverse round trip. +""" + +import math + +import numpy as np +import pytest +from pytest import approx + +import madspace as ms + +BATCH_SIZE = 100000 +M_TOP = 173.0 +M_B = 4.7 +M_W = 80.379 +W_W = 2.085 + +rng = np.random.default_rng(20260811) + + +def massless_volume(n, mass): + """Volume of the n-body massless phase space of a particle of mass ``mass``. + + Phi_n = (2 pi)^(4-3n) (pi/2)^(n-1) M^(2n-4) / ((n-1)! (n-2)!) + """ + return ( + (2 * math.pi) ** (4 - 3 * n) + * (math.pi / 2) ** (n - 1) + * mass ** (2 * n - 4) + / (math.factorial(n - 1) * math.factorial(n - 2)) + ) + + +def two_body_volume(mass, m1, m2): + """Phi_2 = (1/8pi) * sqrt(lambda(M^2, m1^2, m2^2)) / M^2.""" + a, b, c = mass**2, m1**2, m2**2 + lam = a * a + b * b + c * c - 2 * (a * b + a * c + b * c) + return math.sqrt(lam) / (8 * math.pi * a) + + +def decay_mapping(outgoing_masses, propagators, vertices, mass=M_TOP): + diagram = ms.Diagram([mass], outgoing_masses, propagators, vertices) + return ms.PhaseSpaceMapping(ms.Topology(diagram), mass) + + +def sample(mapping): + r = rng.random((BATCH_SIZE, mapping.random_dim())) + p_ext, x1, x2, det = mapping.map_forward([r], []) + return r, p_ext, x1, x2, det + + +def volume(det): + """Monte-Carlo estimate of the volume and its standard error.""" + return np.mean(det), np.std(det) / math.sqrt(len(det)) + + +# -------------------------------------------------------------------------- +# Topology +# -------------------------------------------------------------------------- + +def test_single_incoming_has_no_t_channel(): + diagram = ms.Diagram([M_TOP], [0.0, 0.0], [], [["i0", "o0", "o1"]]) + topology = ms.Topology(diagram) + assert topology.t_propagator_count == 0 + assert topology.incoming_masses == [M_TOP] + assert len(topology.outgoing_masses) == 2 + + +def test_three_or_more_incoming_rejected(): + with pytest.raises(ValueError): + ms.Diagram([1.0, 1.0, 1.0], [0.0, 0.0], [], [["i0", "o0", "o1"]]) + + +# -------------------------------------------------------------------------- +# Kinematics +# -------------------------------------------------------------------------- + +def test_decaying_particle_is_at_rest(): + mapping = decay_mapping([0.0, 0.0], [], [["i0", "o0", "o1"]]) + _, p_ext, _, _, _ = sample(mapping) + # One incoming + two outgoing, not the 2 + n of a collision. + assert p_ext.shape == (BATCH_SIZE, 3, 4) + assert p_ext[:, 0, 0] == approx(M_TOP) + assert p_ext[:, 0, 1:] == approx(0.0, abs=1e-10) + + +def test_momentum_conservation_and_masses(): + mapping = decay_mapping( + [M_B, 0.0, 0.0], + [ms.Propagator(mass=M_W, width=W_W, pdg_id=24)], + [["i0", "o0", "p0"], ["p0", "o1", "o2"]], + ) + _, p_ext, _, _, _ = sample(mapping) + assert p_ext.shape == (BATCH_SIZE, 4, 4) + assert np.sum(p_ext[:, 1:], axis=1) == approx(p_ext[:, 0], abs=1e-8) + + def mass(p): + m2 = p[:, 0] ** 2 - np.sum(p[:, 1:] ** 2, axis=1) + return np.sqrt(np.abs(m2)) + + assert mass(p_ext[:, 0]) == approx(M_TOP, rel=1e-10) + assert mass(p_ext[:, 1]) == approx(M_B, rel=1e-6) + assert mass(p_ext[:, 2]) == approx(0.0, abs=1e-5) + assert mass(p_ext[:, 3]) == approx(0.0, abs=1e-5) + + +def test_no_beam_momentum_fractions(): + mapping = decay_mapping([0.0, 0.0], [], [["i0", "o0", "o1"]]) + _, _, x1, x2, _ = sample(mapping) + assert x1 == approx(1.0) + assert x2 == approx(1.0) + + +def test_random_dim_is_3n_minus_4(): + for n in (2, 3): + mapping = decay_mapping( + [0.0] * n, + [] if n == 2 else [ms.Propagator(mass=M_W, width=W_W, pdg_id=24)], + [["i0", "o0", "o1"]] + if n == 2 + else [["i0", "o0", "p0"], ["p0", "o1", "o2"]], + ) + assert mapping.random_dim() == 3 * n - 4 + assert mapping.particle_count() == n + 1 + + +# -------------------------------------------------------------------------- +# Phase-space volume +# -------------------------------------------------------------------------- + +def test_two_body_massless_volume(): + mapping = decay_mapping([0.0, 0.0], [], [["i0", "o0", "o1"]]) + _, _, _, _, det = sample(mapping) + mean, err = volume(det) + expected = massless_volume(2, M_TOP) + # Two-body is sampled exactly: every weight is the same number. + assert err == approx(0.0, abs=1e-12) + assert mean == approx(expected, rel=1e-10) + + +def test_two_body_massive_volume(): + mapping = decay_mapping([M_B, M_W], [], [["i0", "o0", "o1"]]) + _, _, _, _, det = sample(mapping) + mean, _ = volume(det) + assert mean == approx(two_body_volume(M_TOP, M_B, M_W), rel=1e-10) + + +@pytest.mark.parametrize( + "prop_mass,prop_width,rel_tol", + [ + # A broad propagator makes the Breit-Wigner sampling nearly flat, so the + # estimator has little variance and this is a sharp check of the volume. + (80.0, 60.0, 3e-3), + # The physical W is a narrow resonance: importance sampling it costs + # variance on a flat integrand, so only a loose check is possible here. + (M_W, W_W, 3e-2), + ], + ids=["broad", "narrow-W"], +) +def test_three_body_massless_volume(prop_mass, prop_width, rel_tol): + """t -> b f f' through a Breit-Wigner propagator. The sampled invariant + covers the full range, so whatever the propagator's width, the integral is + the plain massless three-body volume.""" + mapping = decay_mapping( + [0.0, 0.0, 0.0], + [ms.Propagator(mass=prop_mass, width=prop_width, pdg_id=24)], + [["i0", "o0", "p0"], ["p0", "o1", "o2"]], + ) + _, _, _, _, det = sample(mapping) + mean, _ = volume(det) + assert mean == approx(massless_volume(3, M_TOP), rel=rel_tol) + + +# -------------------------------------------------------------------------- +# Invertibility +# -------------------------------------------------------------------------- + +@pytest.mark.parametrize( + "outgoing,propagators,vertices", + [ + ([0.0, 0.0], [], [["i0", "o0", "o1"]]), + ([M_B, M_W], [], [["i0", "o0", "o1"]]), + ( + [M_B, 0.0, 0.0], + [ms.Propagator(mass=M_W, width=W_W, pdg_id=24)], + [["i0", "o0", "p0"], ["p0", "o1", "o2"]], + ), + ], + ids=["2body-massless", "2body-massive", "3body"], +) +def test_forward_inverse_round_trip(outgoing, propagators, vertices): + mapping = decay_mapping(outgoing, propagators, vertices) + r, p_ext, x1, x2, det = sample(mapping) + r_back, det_back = mapping.map_inverse([p_ext, x1, x2], []) + assert r_back == approx(r, abs=1e-8) + assert det * det_back == approx(1.0, rel=1e-8) diff --git a/tests/acceptance_tests/test_cmd.py b/tests/acceptance_tests/test_cmd.py index 5ea141c5c..51bc8c4a4 100755 --- a/tests/acceptance_tests/test_cmd.py +++ b/tests/acceptance_tests/test_cmd.py @@ -4437,6 +4437,97 @@ def has(d): self.assertFalse(has(d), '%s should not survive the regenerate' % d) + def test_output_mg7_decay_subprocess_metadata(self): + """`output mg7` of a decay must describe one initial leg, not two. + + The exporter used to hard-code two initial legs and offset the outgoing + ones by 3, so a 1 -> n process silently produced + incoming = [pdg, None] and lost its first outgoing particle (it landed + on outgoing[-1], overwriting the last one). MadSpin generates its decay + matrix elements exactly this way, so pin the metadata down. + """ + import json + + if os.path.isdir(self.out_dir): + shutil.rmtree(self.out_dir) + + # Without this the leptons come out as merged-particle ids (-82, 83) + # rather than their pdgs, which says nothing about the leg ordering. + # It has to precede the model import, which is what builds the merges. + self.do('set apply_flavor_grouping False') + self.do('import model sm') + self.do('set group_subprocesses False') + self.do('generate t > b w+, w+ > e+ ve') + self.do('output mg7 %s' % self.out_dir) + + with open(pjoin(self.out_dir, 'SubProcesses', + 'subprocesses.json')) as fsock: + subprocesses = json.load(fsock) + self.assertEqual(len(subprocesses), 1) + subproc = subprocesses[0] + + self.assertEqual(subproc['incoming'], [6]) + # The b is the first outgoing leg: it is the one the old offset dropped. + self.assertEqual(subproc['outgoing'], [5, -11, 12]) + # No beam pair, so no beam-swapped mirror configuration. + self.assertFalse(any(flav['mirror'] for flav in subproc['flavors'])) + # The channel topology must hang off the single incoming line i0. + edges = set(edge for channel in subproc['channels'] + for vertex in channel['vertices'] for edge in vertex) + self.assertIn('i0', edges) + self.assertNotIn('i1', edges) + + with open(pjoin(self.out_dir, 'SubProcesses', + 'proc_characteristics')) as fsock: + characteristics = fsock.read() + self.assertIn('ninitial = 1', characteristics) + self.assertIn('nexternal = 4', characteristics) + + def test_output_mg7_decay_run_card_has_no_cuts(self): + """`output mg7` of a decay must ship a run card without any cut. + + A partial width is inclusive, so any kinematic cut biases it low: the + hadron-collider defaults (ptj/ptl/eta/dR) used to survive into a decay + directory and cost ~3% on the t > b w+, w+ > e+ ve width. The card is + written at output time, so the emitted [cuts] section is what has to be + empty -- the user still sees exactly what is run, and can add a cut back + by hand. A collision must keep its defaults untouched. + """ + import madgraph.various.banner as banner_mod + + def cuts_of(path): + card = banner_mod.RunCardMG7(path, consistency=False) + return card['cuts'] + + if os.path.isdir(self.out_dir): + shutil.rmtree(self.out_dir) + + self.do('import model sm') + self.do('generate t > b w+, w+ > e+ ve') + self.do('output mg7 %s' % self.out_dir) + + # both the working card and the "set default" reference + for name in ('run_card.toml', 'run_card_default.toml'): + cuts = cuts_of(pjoin(self.out_dir, 'Cards', name)) + self.assertEqual(dict(cuts), {}, + '%s of a 1 -> n decay must carry no cut, got %s' + % (name, dict(cuts))) + + # the Breit-Wigner cutoff is a sampling range for the off-shell + # propagators, not a cut on the final state: it must survive. + card = banner_mod.RunCardMG7(pjoin(self.out_dir, 'Cards', + 'run_card.toml'), consistency=False) + self.assertEqual(card['phasespace']['bw_cutoff'], 15) + + # a 2 -> n collision keeps the standard cuts + shutil.rmtree(self.out_dir) + self.do('generate p p > t t~') + self.do('output mg7 %s' % self.out_dir) + cuts = cuts_of(pjoin(self.out_dir, 'Cards', 'run_card.toml')) + self.assertEqual(cuts['jet-pt']['min'], 20.0) + self.assertEqual(cuts['lepton-pt']['min'], 10.0) + self.assertEqual(cuts['jet-eta_abs']['max'], 5.0) + @test_manager.bypass_for_py3 def test_madevent_triplet_diquarks(self): """Test MadEvent output of triplet diquarks""" diff --git a/tests/parallel_tests/madspin_benchmark.py b/tests/parallel_tests/madspin_benchmark.py new file mode 100755 index 000000000..5bd41dc88 --- /dev/null +++ b/tests/parallel_tests/madspin_benchmark.py @@ -0,0 +1,412 @@ +#!/usr/bin/env python3 +################################################################################ +# +# 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 +# +################################################################################ +"""Timing benchmark for MadSpin, built on :mod:`madspin_comparator`. + +Runs one production sample through MadSpin once and records the wall-time split +that :class:`~MadSpin.interface_madspin.MadSpinInterface` reports (decay-event +generation, matrix-element generation, max-weight scan, accept/reject loop, +output gzip), plus the optional LHE-parser breakdown. + +The production sample lives under ``--workdir`` and is *reused* across +invocations, so re-timing after a change costs one MadSpin run, not one MG5 run +plus one MadSpin run. + +Usage +----- +Baseline (default benchmark is ``p p > t t~`` with both tops fully decayed):: + + ./tests/parallel_tests/madspin_benchmark.py --label baseline + +Re-time after a change and compare against the baseline record:: + + ./tests/parallel_tests/madspin_benchmark.py --label mg7-decay-gen \\ + --compare-to + +Every run appends a JSON record to ``--out`` (default +``madspin_benchmark_records.json`` inside the workdir). +""" + +from __future__ import absolute_import +from __future__ import division + +import argparse +import datetime +import json +import logging +import os +import platform +import resource +import sys + +pjoin = os.path.join + +_here = os.path.dirname(os.path.realpath(__file__)) +_root = os.path.split(os.path.split(_here)[0])[0] +if _root not in sys.path: + sys.path.insert(0, _root) + +from tests.parallel_tests.madspin_comparator import ( + MadSpinFactory, + SpinModeConfig, +) + +_logger = logging.getLogger('madspin_benchmark') + + +# --------------------------------------------------------------------------- +# Benchmark definitions +# --------------------------------------------------------------------------- +# 'ttbar_full' is the reference benchmark for the MadSpin performance work: +# p p > t t~ with both tops decayed all the way down (t > b w+, w+ > all all), +# which is the case that motivated moving decay-event generation off Fortran +# madevent. +BENCHMARKS = { + 'ttbar_full': dict( + production_process='p p > t t~', + decays=['t > b w+, w+ > all all', + 't~ > b~ w-, w- > all all'], + multiparticles={'p': 'g u d s c u~ d~ s~ c~', + 'j': 'g u d s c u~ d~ s~ c~'}, + extra_run_card={'ebeam1': 6500, 'ebeam2': 6500}, + ), + # Smaller sibling, handy for smoke-testing harness changes without paying + # for the full W decay multiplicity. + 'ttbar_semilep': dict( + production_process='p p > t t~', + decays=['t > b w+, w+ > l+ vl', + 't~ > b~ w-, w- > j j'], + multiparticles={'p': 'g u d s c u~ d~ s~ c~', + 'j': 'g u d s c u~ d~ s~ c~', + 'l+': 'e+ mu+', 'vl': 've vm'}, + extra_run_card={'ebeam1': 6500, 'ebeam2': 6500}, + ), +} + + +def _peak_rss_mb(): + """Peak RSS of this process *and its children*, in MiB. + + ru_maxrss is bytes on macOS and kibibytes on Linux. + """ + usage = resource.getrusage(resource.RUSAGE_CHILDREN) + scale = 1024. * 1024. if sys.platform == 'darwin' else 1024. + return usage.ru_maxrss / scale + + +def run_benchmark(benchmark, label, workdir, nevents, seed, spinmode, + lhe_timers, extra_madspin_settings=None): + """Run one MadSpin timing measurement and return the record dict.""" + spec = BENCHMARKS[benchmark] + # The production sample only depends on (process, nevents, seed), so key + # the working tree on that and let unrelated labels share it. + proc_key = '%s_n%d_s%d' % (benchmark, nevents, seed) + base_dir = pjoin(workdir, proc_key) + if not os.path.isdir(base_dir): + os.makedirs(base_dir) + + factory = MadSpinFactory( + name=proc_key, + nevents=nevents, + seed=seed, + base_dir=base_dir, + extra_madspin_settings=extra_madspin_settings or {}, + **spec + ) + + if lhe_timers: + os.environ['MG_LHE_TIMERS'] = '1' + else: + os.environ.pop('MG_LHE_TIMERS', None) + + config = SpinModeConfig('bench_%s' % label, spinmode) + _logger.info('running benchmark %s [label=%s, spinmode=%s, nevents=%d]', + benchmark, label, spinmode, nevents) + result = factory.run_mode(config) + + record = { + 'label': label, + 'benchmark': benchmark, + 'timestamp': datetime.datetime.now().isoformat(timespec='seconds'), + 'spinmode': spinmode, + 'nevents': nevents, + 'seed': seed, + 'host': platform.node(), + 'platform': platform.platform(), + 'cpu_count': os.cpu_count(), + 'wall_seconds': round(result.wall_seconds, 3), + 'phase_seconds': result.phase_seconds, + 'phase_counts': result.phase_counts, + 'lhe_timers': result.lhe_timers, + 'peak_rss_mb': round(_peak_rss_mb(), 1), + 'efficiency': result.efficiency, + 'BR': result.BR, + 'cross_in': result.cross_in, + 'cross_out': result.cross_out, + 'log_path': result.log_path, + 'lhe_path': result.lhe_path, + } + return record + + +# --------------------------------------------------------------------------- +# Reporting +# --------------------------------------------------------------------------- +# Phases printed in run order rather than alphabetically. Anything MadSpin +# reports that is not listed here is appended afterwards, so a new phase shows +# up without touching this table. +_PHASE_ORDER = [ + 'decay_event_generation', + 'decay_me_generate', + 'decay_me_output', + 'decay_mg7_launch', + 'decay_mg7_integrate', + 'me_generation', + 'max_weight_scan', + 'decay_loop', + 'decay_event_refill', + 'output_gzip', + 'total', +] + +# Phases that run *inside* another phase and are therefore already included in +# its total. Reported indented, and the parent also gets a derived +# " (net)" row so the columns still add up to something meaningful. +_NESTED_IN = { + 'decay_event_refill': 'decay_loop', + # Amplitude generation and code writing happen once, when the decay + # directory is created, so they sit inside the initial generation pass. + 'decay_me_generate': 'decay_event_generation', + 'decay_me_output': 'decay_event_generation', + # ... whereas the launcher runs again for every pool refill, so its time is + # split between decay_event_generation and decay_event_refill and cannot be + # nested under either. + 'decay_mg7_integrate': 'decay_mg7_launch', +} + + +def _ordered_phases(seconds): + known = [p for p in _PHASE_ORDER if p in seconds] + rest = sorted(p for p in seconds if p not in _PHASE_ORDER) + return known + rest + + +# The LHE timers are nested, so they must not simply be added up: +# next_event_readline_total EventFile.next_event, the whole read +# next_event_readline_event_parse the Event() built inside it +# event_parse_total every Event(), including those above +# event_parse_particle_block +# particle_parse_total +# event_parse_tag_block +# event_parse_assign_mother +# Only the two roots count, and event_parse_total must have the part already +# covered by next_event_readline_event_parse taken out of it. +_LHE_NESTED = { + 'next_event_readline_event_parse', + 'event_parse_particle_block', + 'particle_parse_total', + 'event_parse_tag_block', + 'event_parse_assign_mother', +} + + +def lhe_exclusive_seconds(lhe_timers): + """Wall time actually spent in the LHE parser, without double counting.""" + def secs(key): + return lhe_timers.get(key, (0.0, 0))[0] + direct_event_parse = max( + 0.0, secs('event_parse_total') - secs('next_event_readline_event_parse')) + return secs('next_event_readline_total') + direct_event_parse + + +def derived_phases(seconds): + """Return ``seconds`` plus the derived net rows for nested phases.""" + out = dict(seconds) + for child, parent in _NESTED_IN.items(): + if child in seconds and parent in seconds: + out['%s (net)' % parent] = seconds[parent] - seconds[child] + return out + + +def format_record(record, baseline=None): + """Render one record (optionally against a baseline) as a text table.""" + seconds = record.get('phase_seconds', {}) + counts = record.get('phase_counts', {}) + total = seconds.get('total') or record.get('wall_seconds') or 0.0 + base_seconds = (baseline or {}).get('phase_seconds', {}) + + lines = [] + lines.append('MadSpin benchmark: %s [%s, %s, nevents=%s]' + % (record['label'], record['benchmark'], record['spinmode'], + record['nevents'])) + lines.append(' wall %.1f s peak RSS %.0f MiB %s cores' + % (record['wall_seconds'], record['peak_rss_mb'], + record['cpu_count'])) + if baseline: + lines.append(' baseline: %s (wall %.1f s)' + % (baseline['label'], baseline['wall_seconds'])) + + seconds = derived_phases(seconds) + base_seconds = derived_phases(base_seconds) + + header = ' %-28s %10s %7s' % ('phase', 'seconds', '%') + if baseline: + header += ' %10s %8s' % ('base s', 'speedup') + lines.append(header) + for phase in _ordered_phases(seconds): + val = seconds[phase] + share = 100. * val / total if total else 0. + # Nested phases are already counted inside their parent: indent them + # so the shares are not read as additive. + name = (' \\_ ' + phase) if phase in _NESTED_IN else phase + row = ' %-28s %10.2f %6.1f%%' % (name, val, share) + if baseline: + base = base_seconds.get(phase) + if base is None: + row += ' %10s %8s' % ('-', '-') + elif val > 0: + row += ' %10.2f %7.2fx' % (base, base / val) + else: + row += ' %10.2f %8s' % (base, 'inf') + lines.append(row) + + interesting = {k: v for k, v in counts.items() if v} + if interesting: + lines.append(' counts: %s' + % ', '.join('%s=%s' % (k, v) + for k, v in sorted(interesting.items()))) + if record.get('efficiency'): + lines.append(' unweighting efficiency: %.4f' % record['efficiency']) + + lhe = record.get('lhe_timers') or {} + if lhe: + lines.append(' LHE parser: %.2f s (%.1f%% of wall)' + % (lhe_exclusive_seconds(lhe), + 100. * lhe_exclusive_seconds(lhe) / total if total else 0.)) + for key, (secs, calls) in sorted(lhe.items(), key=lambda kv: -kv[1][0]): + marker = ' ' if key in _LHE_NESTED else '* ' + lines.append(' %s%-24s %10.2f s over %d call(s)' + % (marker, key, secs, calls)) + lines.append(' (* counted in the total; the rest are nested ' + 'inside a starred timer)') + return '\n'.join(lines) + + +def load_records(path): + if not os.path.exists(path): + return [] + with open(path) as fp: + try: + return json.load(fp) + except ValueError: + return [] + + +def main(argv=None): + parser = argparse.ArgumentParser(description=__doc__, + formatter_class=argparse.RawDescriptionHelpFormatter) + parser.add_argument('--benchmark', default='ttbar_full', + choices=sorted(BENCHMARKS), + help='which benchmark to run (default: ttbar_full)') + parser.add_argument('--label', default='run', + help='name for this measurement, used in the report') + parser.add_argument('--workdir', default=None, + help='where the production sample and run dirs live ' + '(default: $MADSPIN_BENCH_DIR or ./madspin_bench)') + parser.add_argument('--nevents', type=int, default=10000) + parser.add_argument('--seed', type=int, default=42) + parser.add_argument('--spinmode', default='PA', + help='MadSpin spinmode (default: PA, the shipped default)') + parser.add_argument('--no-lhe-timers', action='store_true', + help='do not set MG_LHE_TIMERS (the timers add a small ' + 'per-particle overhead to the parse path)') + parser.add_argument('--set', dest='settings', action='append', default=[], + metavar='KEY=VALUE', + help='extra "set KEY VALUE" line for the MadSpin card ' + '(repeatable)') + parser.add_argument('--out', default=None, + help='JSON file the records are appended to ' + '(default: /madspin_benchmark_records.json)') + parser.add_argument('--compare-to', default=None, + help='label of a previous record (in --out) to compare ' + 'against; "previous" picks the last one') + parser.add_argument('--report-only', action='store_true', + help='do not run anything, just print the stored records') + args = parser.parse_args(argv) + + logging.basicConfig(level=logging.INFO, format='%(levelname)s: %(message)s') + + workdir = args.workdir or os.environ.get('MADSPIN_BENCH_DIR') \ + or pjoin(os.getcwd(), 'madspin_bench') + workdir = os.path.realpath(workdir) + if not os.path.isdir(workdir): + os.makedirs(workdir) + out_path = args.out or pjoin(workdir, 'madspin_benchmark_records.json') + + records = load_records(out_path) + + if args.report_only: + if not records: + print('no records in %s' % out_path) + return 1 + for rec in records: + print(format_record(rec)) + print('') + return 0 + + extra = {} + for item in args.settings: + if '=' not in item: + parser.error('--set expects KEY=VALUE, got %r' % item) + key, val = item.split('=', 1) + extra[key.strip()] = val.strip() + + record = run_benchmark( + benchmark=args.benchmark, + label=args.label, + workdir=workdir, + nevents=args.nevents, + seed=args.seed, + spinmode=args.spinmode, + lhe_timers=not args.no_lhe_timers, + extra_madspin_settings=extra, + ) + + baseline = None + if args.compare_to: + candidates = [r for r in records + if r['benchmark'] == record['benchmark']] + if args.compare_to == 'previous': + baseline = candidates[-1] if candidates else None + else: + matching = [r for r in candidates if r['label'] == args.compare_to] + baseline = matching[-1] if matching else None + if baseline is None: + _logger.warning('no stored record labelled %r for benchmark %s', + args.compare_to, record['benchmark']) + + records.append(record) + with open(out_path, 'w') as fp: + json.dump(records, fp, indent=2, sort_keys=True) + + print('') + print(format_record(record, baseline)) + print('') + print('record appended to %s' % out_path) + return 0 + + +if __name__ == '__main__': + sys.exit(main()) diff --git a/tests/parallel_tests/madspin_comparator.py b/tests/parallel_tests/madspin_comparator.py index 90552c61c..0b3d206b2 100644 --- a/tests/parallel_tests/madspin_comparator.py +++ b/tests/parallel_tests/madspin_comparator.py @@ -28,6 +28,7 @@ from __future__ import division import collections +import json import logging import math import os @@ -106,11 +107,18 @@ class MadSpinResult(object): def __init__(self, config, lhe_path, log_path, wall_seconds, BR, BR_err, efficiency, nevents_in, - cross_out=None, cross_in=None): + cross_out=None, cross_in=None, + phase_seconds=None, phase_counts=None, lhe_timers=None): self.config = config self.lhe_path = lhe_path self.log_path = log_path self.wall_seconds = wall_seconds + # Per-phase wall time / occurrence counts as reported by MadSpin + # itself (see MadSpinInterface._log_phase_timings), plus the optional + # LHE-parser breakdown when MG_LHE_TIMERS was set for the run. + self.phase_seconds = dict(phase_seconds or {}) + self.phase_counts = dict(phase_counts or {}) + self.lhe_timers = dict(lhe_timers or {}) self.BR = BR self.BR_err = BR_err self.efficiency = efficiency @@ -170,6 +178,44 @@ def count_pdgs(self): _RE_WRITTEN = re.compile( r'Total number of events written:\s*(\d+)\s*/\s*(\d+)' ) +# MadSpin logs through madgraph's ColorFormatter, which emits ANSI escapes +# unconditionally -- also into a redirected file. Strip them before matching. +_RE_ANSI = re.compile(r'\x1b\[[0-9;]*m') +# Machine-readable per-phase wall times emitted by +# MadSpinInterface._log_phase_timings at the end of every run. The greedy +# ``.*`` stops at the last '}' on the line, i.e. the end of the JSON payload, +# so trailing formatter decoration does not matter. +_RE_PHASES = re.compile(r'MadSpin phase timings:\s*(\{.*\})') +# Optional LHE parser timers (only present when MG_LHE_TIMERS is set): +# " Event.__init__: 1.234567s total over 42 call(s) (avg 0.029394s)" +_RE_LHE_TIMER = re.compile( + r'^\s+(\S+):\s*([0-9.eE+-]+)s total over (\d+) call\(s\)', re.MULTILINE +) + + +def parse_phase_timings(text): + """Return ``(seconds, counts)`` dicts from a MadSpin log, or ``({}, {})``. + + Uses the last occurrence so a log holding several runs reports the last one. + """ + payload = None + for match in _RE_PHASES.finditer(_RE_ANSI.sub('', text)): + payload = match.group(1) + if payload is None: + return {}, {} + try: + data = json.loads(payload) + except ValueError: + return {}, {} + return data.get('seconds', {}), data.get('counts', {}) + + +def parse_lhe_timers(text): + """Return ``{key: (seconds, calls)}`` from the LHE parser timing summary.""" + out = {} + for match in _RE_LHE_TIMER.finditer(_RE_ANSI.sub('', text)): + out[match.group(1)] = (float(match.group(2)), int(match.group(3))) + return out def _parse_log(text): @@ -285,25 +331,53 @@ def _write_mg5_script(self, script_path): # plain unweighted_events.lhe here). lines.append('output madevent %s' % self.proc_dir) lines.append('launch %s' % self.proc_dir) + # ``launch`` opens ONE menu that takes both the tool switches and the + # ``set `` lines; ``done`` closes it and starts the run. + # Putting the set lines after a first ``done`` (as this used to) meant + # they were never executed -- the run had already started, so nevents, + # iseed and the beam energies silently kept their run_card defaults. lines.append('madspin=OFF') # MadSpin runs separately, mode by mode lines.append('shower=OFF') lines.append('detector=OFF') lines.append('analysis=OFF') - lines.append('done') # end card edit menu lines.append('set nevents %d' % self.nevents) lines.append('set iseed %d' % self.seed) + # Systematics needs lhapdf's python bindings, which are an optional + # (and, on some interpreters, broken) dependency; MadSpin does not use + # the reweighting information, so keep the whole step out of the run. lines.append('set use_syst False') + lines.append('set systematics_program none') for key, val in self.extra_run_card.items(): lines.append('set %s %s' % (key, val)) - lines.append('done') # end second card edit menu (after card adjustments) + lines.append('done') # close the menu -> start the run with open(script_path, 'w') as fp: fp.write('\n'.join(lines) + '\n') + def _existing_production(self): + """Return an already-generated production LHE under ``proc_dir``, if any. + + Only reachable when the caller passed an explicit ``base_dir`` (the + default tempdir is fresh every time). This is what lets a benchmark + re-time MadSpin repeatedly without paying for the production run again. + """ + for name in ('unweighted_events.lhe.gz', 'unweighted_events.lhe'): + candidate = pjoin(self.proc_dir, 'Events', 'run_01', name) + if os.path.exists(candidate): + return candidate + return None + def produce_events(self): """Run mg5_aMC once; cache the LHE file path.""" if self.events_file: return self.events_file + cached = self._existing_production() + if cached: + _logger.info('%s: reusing production sample %s', self.name, cached) + self.events_file = cached + self.cross_in = _read_lhe_cross(self.events_file) + return self.events_file + script_path = pjoin(self.base_dir, 'mg5_script.dat') self._write_mg5_script(script_path) @@ -329,7 +403,11 @@ def produce_events(self): log_text = fp.read() for marker in ('NoDiagramException', 'command not executed: output', - 'command not executed: launch'): + 'command not executed: launch', + # A command that raised leaves the rest of the script + # unexecuted, so the run_card 'set' lines silently do not + # apply and the sample is not the one that was asked for. + 'interrupted with error'): if marker in log_text: raise RuntimeError( 'mg5_aMC aborted mid-script for factory %s ' @@ -449,6 +527,8 @@ def run_mode(self, config): BR, accepted, trials, efficiency = _parse_log(log_text) if efficiency is None and accepted is not None and trials: efficiency = float(accepted) / float(trials) + phase_seconds, phase_counts = parse_phase_timings(log_text) + lhe_timers = parse_lhe_timers(log_text) # Always read the decayed banner's cross-section -- this is the # physics-observable we want to compare across modes. @@ -476,6 +556,9 @@ def run_mode(self, config): nevents_in=self.nevents, cross_out=cross_out, cross_in=getattr(self, 'cross_in', None), + phase_seconds=phase_seconds, + phase_counts=phase_counts, + lhe_timers=lhe_timers, ) self._results[config.label] = result return result