diff --git a/madgraph/interface/madgraph_interface.py b/madgraph/interface/madgraph_interface.py index 836112fc28..ce76124ce1 100755 --- a/madgraph/interface/madgraph_interface.py +++ b/madgraph/interface/madgraph_interface.py @@ -178,6 +178,26 @@ def interference_mode_of(option, orders): return '' +# '[amplicol]' takes the matrix elements of a process from AmpliCol, a +# colour-ordered generator, instead of MG diagrams (mg7 output only) +ME_BACKENDS = ('amplicol',) +_BACKEND_BRACKET = re.compile(r"\[\s*(?P\w+)\s*\]") + + +def split_backend_line(line): + """(line without the bracket, backend) for a process with a matrix-element + backend bracket like '[amplicol]', (line, '') for any other line.""" + + match = _BACKEND_BRACKET.search(line) + if not match or match.group('backend').lower() not in ME_BACKENDS: + return line, '' + rest = line[:match.start()] + ' ' + line[match.end():] + if '[' in rest or ']' in rest: + raise madgraph.InvalidCmd( + "[%s] cannot be combined with another '[...]'." % match.group('backend')) + return ' '.join(rest.split()), match.group('backend').lower() + + def split_interference_line(line): """Split a process definition of the form @@ -2345,7 +2365,11 @@ def check_output(self, args, default='mg7'): return - if not self._curr_amps: + if self._curr_amplicol_defs and self._export_format != 'mg7': + raise self.InvalidCmd( + 'The [amplicol] processes are only supported by the mg7 output.') + + if not self._curr_amps and not self._curr_amplicol_defs: text = 'No processes generated. Please generate a process first.' raise self.InvalidCmd(text) @@ -3866,6 +3890,8 @@ def _tutorial_opts(self): _curr_model = None #base_objects.Model() _curr_amps = diagram_generation.AmplitudeList() _curr_proc_defs = base_objects.ProcessDefinitionList() + # processes whose matrix elements come from AmpliCol, see split_backend_line + _curr_amplicol_defs = base_objects.ProcessDefinitionList() _curr_matrix_elements = helas_objects.HelasMultiProcess() _curr_helas_model = None _curr_exporter = None @@ -3954,6 +3980,7 @@ def setup(self): self._done_export=False self._curr_amps = diagram_generation.AmplitudeList() self._curr_proc_defs = base_objects.ProcessDefinitionList() + self._curr_amplicol_defs = base_objects.ProcessDefinitionList() self._curr_matrix_elements = helas_objects.HelasMultiProcess() self._v4_export_formats = ['madevent', 'standalone_fortran','standalone_msP','standalone_msF', @@ -3978,6 +4005,30 @@ def do_quit(self, line): # Add a process to the existing multiprocess definition # Generate a new amplitude + def add_backend_process(self, procdef, backend): + """Keep a process whose matrix elements come from another backend than + MG diagrams. It is exported by the mg7 output, without diagrams.""" + + if procdef.get('decay_chains') or procdef.get('required_s_channels') or \ + procdef.get('forbidden_onsh_s_channels') or \ + procdef.get('forbidden_s_channels') or \ + procdef.get('forbidden_particles') or procdef.get('orders') or \ + procdef.get('squared_orders') or procdef.get('is_decay_chain') or \ + any(leg.get('polarization') for leg in procdef.get('legs')) or \ + procdef.get_ninitial() != 2: + raise self.InvalidCmd( + "[%s] supports 2 -> n processes without decay chains, coupling " + "orders, s-channel or particle restrictions and polarizations." + % backend) + initial = [amp.get_ninitial() for amp in self._curr_amps] + if initial and initial[0] != procdef.get_ninitial(): + raise self.InvalidCmd( + "Can not mix processes with different number of initial states.") + self._curr_amplicol_defs.append(procdef) + self._done_export = False + logger.info("Process %s is computed by %s at output time" % + (procdef.nice_string(prefix=False), backend)) + def do_add(self, line, counter=0): """Generate an amplitude for a given process and add to existing amplitudes @@ -4079,7 +4130,10 @@ def do_add(self, line, counter=0): else: nb_proc = len([l for l in self.history if l.startswith(('generate','add process'))]) nb_proc += counter + line, backend = split_backend_line(line) myprocdef = self.extract_process(line, proc_number=nb_proc) + if backend: + return self.add_backend_process(myprocdef, backend) # 'left [treextree] right': hidden order selecting the # interference, minimal orders of the right-hand process diagram_generation.prepare_interference_process(myprocdef) @@ -6045,6 +6099,7 @@ def clean_process(self): self._curr_amps = diagram_generation.AmplitudeList() # Reset Process definition self._curr_proc_defs = base_objects.ProcessDefinitionList() + self._curr_amplicol_defs = base_objects.ProcessDefinitionList() # Reset Helas matrix elements self._curr_matrix_elements = helas_objects.HelasMultiProcess() self._generate_info = "" @@ -7555,6 +7610,7 @@ def do_import(self, line, force=False, options={}): self._curr_amps = diagram_generation.AmplitudeList() # Reset proc defs self._curr_proc_defs = base_objects.ProcessDefinitionList() + self._curr_amplicol_defs = base_objects.ProcessDefinitionList() self._curr_matrix_elements = helas_objects.HelasMultiProcess() process_checks.store_aloha = [] self.advise_neglected_masses() @@ -10938,6 +10994,7 @@ def set2_group_subprocesses(self, args, log=True): logger.info('Note that you need to regenerate all processes') self._curr_amps = diagram_generation.AmplitudeList() self._curr_proc_defs = base_objects.ProcessDefinitionList() + self._curr_amplicol_defs = base_objects.ProcessDefinitionList() self._curr_matrix_elements = helas_objects.HelasMultiProcess() if self.options[args[0]] == 'gpu': export_v4.ProcessExporterFortranMEGroup.grouped_mode = 'gpu' @@ -11113,6 +11170,7 @@ def set2_gauge(self, args, log=True): self._curr_model = None self._curr_amps = diagram_generation.AmplitudeList() self._curr_proc_defs = base_objects.ProcessDefinitionList() + self._curr_amplicol_defs = base_objects.ProcessDefinitionList() self._curr_matrix_elements = helas_objects.HelasMultiProcess() self._curr_helas_model = None self._curr_exporter = None @@ -11592,6 +11650,7 @@ def set2_OLP(self, args, log=True): # depend on this option self._curr_amps = diagram_generation.AmplitudeList() self._curr_proc_defs = base_objects.ProcessDefinitionList() + self._curr_amplicol_defs = base_objects.ProcessDefinitionList() self._curr_matrix_elements = helas_objects.HelasMultiProcess() self._curr_exporter = None self._second_exporter = None @@ -12485,7 +12544,7 @@ def do_output(self, line): # have been defined with multiparticle labels, because then # branching ratios necessitates to keep subprocesses independent. # That applies only if there is more than one subprocess of course. - if self._curr_amps[0].get_ninitial() == 1 and \ + if self._curr_amps and self._curr_amps[0].get_ninitial() == 1 and \ len(self._curr_amps)>1: processes = [amp.get('process') for amp in self._curr_amps if 'process' in list(amp.keys())] @@ -12881,7 +12940,7 @@ def generate_matrix_elements(self, group_processes=True): matrix_elements = self._curr_matrix_elements.get_matrix_elements() - if any(l.get('flavor') is None for l in matrix_elements[0].get('processes')[0].get('legs')): + if matrix_elements and any(l.get('flavor') is None for l in matrix_elements[0].get('processes')[0].get('legs')): [m.get_external_flavors() for m in matrix_elements] #precompute all flavors # Just the matrix.f files if self._export_format == 'matrix': @@ -13544,6 +13603,7 @@ def do_decay_diagram(self, line, skip_2body=False, model=None): # Reset amplitudes self._curr_amps = diagram_generation.AmplitudeList() self._curr_proc_defs = base_objects.ProcessDefinitionList() + self._curr_amplicol_defs = base_objects.ProcessDefinitionList() # Reset Helas matrix elements self._curr_matrix_elements = helas_objects.HelasMultiProcess() # Reset _done_export, since we have new process diff --git a/madgraph/interface/master_interface.py b/madgraph/interface/master_interface.py index 7db93f4e83..dd440419a9 100755 --- a/madgraph/interface/master_interface.py +++ b/madgraph/interface/master_interface.py @@ -173,6 +173,10 @@ def extract_process_type(line): space_before = re.compile(r"(?P\S)(?P[\\[\\]/\,\\$\\>|])(?P\S)") line2 = space_before.sub(r'\g \g \g', line) + # a matrix-element backend bracket like [amplicol] is a tree process + if MGcmd.split_backend_line(re.split(r'\s\-\-', line, 1)[0])[1]: + return ('tree', None, []) + # Use regular expressions to extract the loop mode (if present) and its # option, specified in the line with format [ option = loop_orders ] or # [ loop_orders ] which implicitly select the 'all' option. diff --git a/madgraph/iolibs/amplicol_worker.py b/madgraph/iolibs/amplicol_worker.py new file mode 100644 index 0000000000..74001176c6 --- /dev/null +++ b/madgraph/iolibs/amplicol_worker.py @@ -0,0 +1,81 @@ +################################################################################ +# +# Copyright (c) 2026 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 +# +################################################################################ +"""Generation of the '[amplicol]' artifacts with pyamplicol, run in its own +process by export_amplicol. + +Symbolica, which pyamplicol imports, holds a per-user lock until the process +exits, so it is never imported by MadGraph itself: sessions do not block each +other while they are open. Usage: python amplicol_worker.py JOB.json +""" + +import json +import logging +import sys + + +def model_source(pyamplicol, model): + """The model compiled for pyamplicol, cached in the output.""" + import os + if not os.path.exists(model['cache']): + from pyamplicol.config import EvaluatorConfig, JITConfig + source = pyamplicol.ModelSource.from_path(model['path'], + restriction=model['restriction']) + source.compile(prepared_output=model['cache'], + evaluator=EvaluatorConfig(jit=JITConfig(optimization_level=2))) + return pyamplicol.ModelSource.from_path(model['cache']) + + +def generate_artifact(pyamplicol, artifact, accuracy, process, model): + from pyamplicol.config import resolve_config + # workers and cores are left to pyamplicol, which clamps them without a license + config = resolve_config({ + 'schema_version': 1, + 'action': 'generate', + 'color': {'accuracy': accuracy}, + 'process': {'multiparticles': process['multiparticles'], + 'coupling_order_policy': 'explicit', + 'max_coupling_orders': process['orders']}, + 'generation': {'emit_api_bundle': True, + 'relation_discovery': {'mode': 'off'}}, + 'evaluator': {'backend': 'jit', 'execution_mode': 'recurrence', + 'jit': {'optimization_level': 2}}, + }) + pyamplicol.Generator(config).generate(process['expression'], artifact, + model=model, mode='replace') + + +def main(job_path): + import os + logging.basicConfig(level=logging.WARNING, format='%(name)s: %(message)s') + with open(job_path) as stream: + job = json.load(stream) + import pyamplicol + from pyamplicol.artifacts import umami_reweight + model = model_source(pyamplicol, job['model']) + providers = [] + for process in job['processes']: + path = process['path'] + for accuracy in ('lc', 'full'): + generate_artifact(pyamplicol, os.path.join(path, accuracy), accuracy, + process, model) + providers.append(umami_reweight.write_reweight_provider( + os.path.join(path, 'lc'), os.path.join(path, 'full'), + os.path.join(path, 'reweight'))) + with open(job['result'], 'w') as stream: + json.dump(providers, stream) + + +if __name__ == '__main__': + main(sys.argv[1]) diff --git a/madgraph/iolibs/export_amplicol.py b/madgraph/iolibs/export_amplicol.py new file mode 100644 index 0000000000..c278c40453 --- /dev/null +++ b/madgraph/iolibs/export_amplicol.py @@ -0,0 +1,362 @@ +################################################################################ +# +# Copyright (c) 2026 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 +# +################################################################################ +"""mg7 output of the processes computed by AmpliCol ('[amplicol]'). + +pyamplicol generates the colour-ordered matrix elements of a process definition +in SubProcesses/A_: a leading-colour artifact (lc), a full-colour one +(full), and the sources of the UMAMI library returning their ratio (reweight). +Its makefile builds the libraries into lib/. Every colour ordering (UMAMI +channel) of every provider becomes one entry of subprocesses.json, integrated +on its own with a colour-ordered phase-space mapping. +""" + +import importlib.util +import json +import logging +import os +import re +import shutil +import subprocess +import sys +import tempfile +import time + +import madgraph + +logger = logging.getLogger('madgraph.export_amplicol') +pjoin = os.path.join + +MODEL_FILE = 'amplicol_model.pyamplicol-model' + +MAKEFILE = """\ +# Builds the AmpliCol matrix elements of this process into ../../lib. +RUSTICOL_CONFIG ?= %(rusticol_config)s +BUILD := $(CURDIR)/build + +all: +\t$(MAKE) -C lc/API/umami shared BUILD_DIR=$(BUILD)/lc RUSTICOL_CONFIG=$(RUSTICOL_CONFIG) +\t$(MAKE) -C reweight shared BUILD_DIR=$(BUILD)/reweight RUSTICOL_CONFIG=$(RUSTICOL_CONFIG) +%(copies)s +clean: +\trm -rf $(BUILD) +""" + + +WORKER = pjoin(os.path.dirname(os.path.abspath(__file__)), 'amplicol_worker.py') +# printed by Symbolica when another process holds its per-user thread +SYMBOLICA_BUSY = 'Cannot start another restricted Symbolica thread' + + +def check_pyamplicol(): + """pyamplicol is only imported by the worker, see amplicol_worker.""" + if importlib.util.find_spec('pyamplicol') is None: + raise madgraph.InvalidCmd( + 'The [amplicol] processes need pyamplicol. Install it in the Python ' + 'environment of MadGraph.') + + +def run_worker(job): + """Run amplicol_worker on a job, waiting while Symbolica is busy elsewhere.""" + env = dict(os.environ, SYMBOLICA_HIDE_BANNER='1', + PYTHONPATH=os.pathsep.join(p for p in sys.path if p)) + waiting = False + while True: + worker = subprocess.Popen([sys.executable, WORKER, job], env=env, + stdout=subprocess.PIPE, stderr=subprocess.STDOUT, + text=True) + busy = False + for line in worker.stdout: + if SYMBOLICA_BUSY in line: + busy = True + elif not line.startswith(('\u250c', '\u2502', '\u2514', '\u251c')): + sys.stdout.write(line) + if worker.wait() == 0: + return + if not busy: + raise madgraph.InvalidCmd('pyamplicol failed to generate the [amplicol] ' + 'matrix elements, see the messages above') + if 'symbolica' in sys.modules: + raise madgraph.InvalidCmd( + 'Symbolica is imported in this MadGraph process, so the pyamplicol ' + 'worker cannot start: do not import pyamplicol or symbolica here') + if not waiting: + logger.info('Symbolica is used by another process: waiting for it') + waiting = True + time.sleep(10) + + +def load_metadata(artifact): + """The UMAMI metadata documents of an artifact, one per provider.""" + root = pjoin(artifact, 'API', 'umami') + with open(pjoin(root, 'metadata.json')) as stream: + data = json.load(stream) + if data.get('kind') != 'pyamplicol-umami-provider-index': + return [data] + documents = [] + for provider in data['providers']: + with open(pjoin(root, provider['metadata'])) as stream: + documents.append(json.load(stream)) + return documents + + +def leg_pdgs(leg, model): + """The pdg ids of a leg: the flavours MG merged into a group id, or the ids + with the group ids expanded.""" + if leg.get('flavor'): + return list(leg.get('flavor')) + merged = model.get('merged_particles') + pdgs = [] + for pid in leg.get('ids'): + members = merged.get(abs(pid), [abs(pid)]) + pdgs.extend(member if pid > 0 else -member for member in members) + return pdgs + + +def process_name(procdef, model, multiparticles): + """A directory name for a process definition, e.g. 'pp_jjj', with the + multiparticle labels of the interface.""" + particles = model.get('particle_dict') + labels = [] + for leg in procdef.get('legs'): + pdgs = leg_pdgs(leg, model) + matches = [label for label, ids in multiparticles.items() + if sorted(ids) == sorted(pdgs)] + preferred = 'j' if leg.get('state') else 'p' + label = preferred if preferred in matches else next(iter(matches), None) + if label is None: + label = ''.join(particles[pid].get_name() for pid in pdgs) + labels.append((leg.get('state'), label)) + name = ''.join(l for state, l in labels if not state) + '_' + \ + ''.join(l for state, l in labels if state) + return re.sub(r'[^A-Za-z0-9_]', '', + name.replace('~', 'x').replace('+', 'p').replace('-', 'm')) + + +def process_expression(procdef, model): + """The process in pyamplicol syntax, and its multiparticle definitions.""" + particles = model.get('particle_dict') + multiparticles = {} + labels = [] + for leg in procdef.get('legs'): + names = [particles[pid].get_name() for pid in leg_pdgs(leg, model)] + if len(names) == 1: + labels.append(names[0]) + continue + label = next((key for key, value in multiparticles.items() + if value == names), None) + if label is None: + label = 'mp%d' % len(multiparticles) + multiparticles[label] = names + labels.append(label) + initial = [label for leg, label in zip(procdef.get('legs'), labels) + if not leg.get('state')] + final = [label for leg, label in zip(procdef.get('legs'), labels) + if leg.get('state')] + return '%s > %s' % (' '.join(initial), ' '.join(final)), multiparticles + + +def coupling_orders(procdef, model): + """The largest orders of the couplings above the lowest hierarchy, as MG + selects them by default: the minimal WEIGHTED order of the definition, + with n - 2 vertices at tree level.""" + import madgraph.core.diagram_generation as diagram_generation + weighted = diagram_generation.MultiProcess.find_optimal_process_orders( + procdef)['WEIGHTED'] + hierarchy = model.get('order_hierarchy') + lowest = min(hierarchy.values()) + vertices = len(procdef.get('legs')) - 2 + return {order: max((weighted - vertices * lowest) // (weight - lowest), 0) + for order, weight in hierarchy.items() if weight > lowest} + + +def model_job(model, cache_path): + """The UFO model of the output with its restriction, compiled for + pyamplicol once per output.""" + path = model.get('modelpath') + reference = os.path.basename(model.get('modelpath+restriction')) + restriction = 'default' + if '-' in reference: + suffix = reference.split('-', 1)[1] + if suffix == 'full': + restriction = 'none' + else: + restriction = pjoin(path, 'restrict_%s.dat' % suffix) + if not os.path.exists(cache_path): + logger.info('Compiling the model %s for AmpliCol' % reference) + return {'path': path, 'restriction': restriction, 'cache': cache_path} + + +def representative_processes(proc_defs, model): + """Processes standing for the [amplicol] definitions where the output needs + MG processes (run card defaults, proc_characteristics): every initial state, + with the first particle of each final-state leg.""" + import itertools + import madgraph.core.base_objects as base_objects + processes = [] + for procdef in proc_defs: + legs = procdef.get('legs') + initial = [leg_pdgs(leg, model) for leg in legs if not leg.get('state')] + final = [leg_pdgs(leg, model)[0] for leg in legs if leg.get('state')] + for ids in itertools.product(*initial): + pdgs = list(ids) + final + processes.append(base_objects.Process({ + 'legs': base_objects.LegList([ + base_objects.Leg({'id': pdg, 'number': number + 1, + 'state': number >= len(ids)}) + for number, pdg in enumerate(pdgs)]), + 'model': model, + 'id': procdef.get('id'), + })) + return processes + + +def color_types(model, pdgs): + types = {} + for pdg in pdgs: + particle = model.get_particle(pdg) + types[pdg] = particle.get_color() + types[-pdg] = model.get_particle(particle.get_anti_pdg_code()).get_color() + return types + + +def mirrored(pdgs, incoming, initial_legs, exported): + """Whether the process with swapped beams is requested but not exported: + pyamplicol keeps one process per unordered initial state.""" + if incoming != 2 or pdgs[0] == pdgs[1]: + return False + swapped = ((pdgs[1], pdgs[0]), tuple(sorted(pdgs[2:]))) + return (pdgs[1] in initial_legs[0] and pdgs[0] in initial_legs[1] + and swapped not in exported) + + +def subprocess_entries(metadata, model, rel_dir, me_path, ratio_path, qcd_power, + initial_legs): + """One subprocesses.json entry per channel of a UMAMI provider; initial_legs + are the pdg ids allowed for each initial leg.""" + pdg_ids = metadata['pdg_ids'] + provider = metadata['provider'] + incoming = provider['incoming_count'] + exported = {(tuple(pdgs[:incoming]), tuple(sorted(pdgs[incoming:]))) + for pdgs in pdg_ids} + entries = [] + for channel_index, channel in enumerate(metadata['channels']): + flavors = [] + for flavor_index, process in enumerate(channel['processes']): + if process['multichannels'] != [channel_index]: + raise madgraph.MadGraph5Error( + 'AmpliCol contributions shared by several channels are not ' + 'supported yet') + for term in process['matrix_elements']: + pdgs = pdg_ids[term['pdg_ids']] + flavors.append({ + 'index': flavor_index, + 'options': [pdgs], + 'factor': term['factor'], + 'mirror': mirrored(pdgs, incoming, initial_legs, exported), + }) + first = flavors[0]['options'][0] + entries.append({ + 'backend': 'amplicol', + 'incoming': first[:incoming], + 'outgoing': first[incoming:], + 'channels': [], + 'path': rel_dir, + 'me_path': me_path, + 'ratio_path': ratio_path, + 'artifact': pjoin(rel_dir, 'lc'), + 'ratio_artifact': pjoin(rel_dir, 'reweight'), + 'umami_channel': channel_index, + 'phasespace_order': channel['phasespace_order'], + 'masses': provider['masses'], + 'supports_alpha_s': provider['supports_alpha_s'], + 'flavors': flavors, + 'qcd_power': qcd_power, + 'color_flows': metadata['color_flows'], + 'pdg_color_types': color_types( + model, {abs(pdg) for pdgs in pdg_ids for pdg in pdgs}), + 'helicities': metadata['helicities'], + 'diagram_count': 0, + 'interference': False, + 'parameters': parameter_references(metadata), + }) + return entries + + +def parameter_references(metadata): + """The parameter card entry of every mutable runtime parameter.""" + references = {} + for process in metadata['runtime_processes']: + for parameter in process['parameters']: + if parameter.get('mutable') and parameter.get('lhablock'): + references[parameter['name']] = [parameter['lhablock'], + parameter['lhacode']] + return references + + +def export_amplicol_processes(dir_path, proc_defs, model, multiparticles=None): + """Write SubProcesses/A_ for every '[amplicol]' process definition + and return the subprocesses.json entries.""" + check_pyamplicol() + subproc_dir = pjoin(dir_path, 'SubProcesses') + rusticol_config = shutil.which('rusticol-config') or 'rusticol-config' + processes = [] + for number, procdef in enumerate(proc_defs, start=1): + name = 'A%d_%s' % (number, process_name(procdef, model, multiparticles or {})) + path = pjoin(subproc_dir, name) + expression, process_multiparticles = process_expression(procdef, model) + os.makedirs(path, exist_ok=True) + logger.info('Generating the AmpliCol matrix elements of %s in %s' % + (procdef.nice_string(prefix=False), path)) + processes.append({'name': name, 'path': path, 'expression': expression, + 'multiparticles': process_multiparticles, + 'orders': coupling_orders(procdef, model)}) + with tempfile.TemporaryDirectory(prefix='amplicol_job_') as scratch: + job = pjoin(scratch, 'job.json') + with open(job, 'w') as stream: + json.dump({'model': model_job(model, pjoin(subproc_dir, MODEL_FILE)), + 'processes': processes, + 'result': pjoin(scratch, 'providers.json')}, stream) + run_worker(job) + with open(pjoin(scratch, 'providers.json')) as stream: + providers = json.load(stream) + entries = [] + for procdef, process, provider_ids in zip(proc_defs, processes, providers): + name, path, orders = process['name'], process['path'], process['orders'] + rel_dir = pjoin('SubProcesses', name) + # every vertex carries one coupling: the rest are powers of g_s + qcd_power = (len(procdef.get('legs')) - 2 - sum(orders.values()) + if len(orders) == 1 else -1) + initial_legs = [set(leg_pdgs(leg, model)) + for leg in procdef.get('legs') if not leg.get('state')] + copies = [] + for provider_id, metadata in zip(provider_ids, load_metadata(pjoin(path, 'lc'))): + library = metadata['provider']['library'] + me_path = pjoin('lib', 'libamplicol_%s_%s.so' % (name, provider_id)) + ratio_path = pjoin('lib', 'libamplicol_%s_%s_ratio.so' % + (name, provider_id)) + copies.append('\tcp $(BUILD)/lc/%s ../../%s' % (library, me_path)) + copies.append('\tcp $(BUILD)/reweight/libreweight_%s.so ../../%s' % + (provider_id, ratio_path)) + entries.extend(subprocess_entries(metadata, model, rel_dir, me_path, + ratio_path, qcd_power, initial_legs)) + with open(pjoin(path, 'makefile'), 'w') as stream: + stream.write(MAKEFILE % {'rusticol_config': rusticol_config, + 'copies': '\n'.join(copies)}) + with open(pjoin(path, 'process.json'), 'w') as stream: + json.dump({'process': procdef.nice_string(prefix=False), + 'expression': process['expression'], + 'multiparticles': process['multiparticles'], + 'max_coupling_orders': orders}, stream, indent=2) + return entries diff --git a/madgraph/iolibs/export_cpp.py b/madgraph/iolibs/export_cpp.py index e70cbfe98c..e5d766e9c0 100755 --- a/madgraph/iolibs/export_cpp.py +++ b/madgraph/iolibs/export_cpp.py @@ -3458,6 +3458,7 @@ def get_merged_info(self): return merged_subproc_info def finalize(self, matrix_elements=None, history='', *args, **kwargs): + self.export_amplicol() file_name = os.path.normpath(os.path.join( self.dir_path, "SubProcesses", "subprocesses.json" )) @@ -3504,8 +3505,35 @@ def finalize(self, matrix_elements=None, history='', *args, **kwargs): def pass_information_from_cmd(self, cmd): """Capture the process definitions from the command interface; needed to - generate the MadAnalysis5 default cards at output time.""" + generate the MadAnalysis5 default cards at output time, and the + [amplicol] processes.""" self.proc_defs = getattr(cmd, '_curr_proc_defs', None) + self.amplicol_defs = getattr(cmd, '_curr_amplicol_defs', None) or [] + # without MG diagrams there is no madmatrix code to write or build + self.has_madmatrix = bool(cmd._curr_amps) or not self.amplicol_defs + self.amplicol_model = cmd._curr_model + self.amplicol_multiparticles = dict(getattr(cmd, '_multiparticles', {})) + + def amplicol_processes(self): + """MG processes standing for the [amplicol] processes.""" + if not getattr(self, 'amplicol_defs', None): + return [] + import madgraph.iolibs.export_amplicol as export_amplicol + return export_amplicol.representative_processes( + self.amplicol_defs, self.amplicol_model) + + def export_amplicol(self): + """The [amplicol] processes, as subprocesses after the madmatrix ones.""" + if not getattr(self, 'amplicol_defs', None): + return + import madgraph.iolibs.export_amplicol as export_amplicol + for entry in export_amplicol.export_amplicol_processes( + self.dir_path, self.amplicol_defs, self.amplicol_model, + self.amplicol_multiparticles): + # never merged with another subprocess + self.merged_subprocesses[('amplicol', len(self.process_info))].append( + (len(self.process_info), [])) + self.process_info.append(entry) def create_ma5_default_cards(self, matrix_elements, history): """Call MadAnalysis5 to write process-tailored default analysis cards @@ -3595,6 +3623,8 @@ def create_run_card(self, matrix_elements, history): for me in matrix_elements['matrix_elements']] except (KeyError, TypeError): processes = None + if self.amplicol_processes(): + processes = (processes or []) + [self.amplicol_processes()] if processes: run_card.create_default_for_process(self.proc_characteristic, @@ -3631,8 +3661,12 @@ def create_proc_characteristics(self, matrix_elements): else: me_list = [] - procs = [] - qcd_orders = set() + procs = list(self.amplicol_processes()) + qcd_orders = set(entry['qcd_power'] for entry in self.process_info + if entry.get('backend') == 'amplicol') + if procs: + pc['nexternal'] = max(pc['nexternal'], len(procs[0].get('legs'))) + pc['ninitial'] = 2 for me in me_list: if not me.get('processes'): continue diff --git a/madgraph/iolibs/template_files/mg7/cli.py b/madgraph/iolibs/template_files/mg7/cli.py new file mode 100644 index 0000000000..d096cca0d3 --- /dev/null +++ b/madgraph/iolibs/template_files/mg7/cli.py @@ -0,0 +1,470 @@ +"""The MG7> command line interface and its card editing.""" + +import argparse +import os +import re +import resource + +from madgraph.iolibs.template_files.mg7.common import logger +from madgraph.interface.extended_cmd import Cmd +from madgraph.various.banner import RunCardMG7 +from madgraph.iolibs.template_files.mg7 import npy_to_lhe +from madgraph.iolibs.template_files.mg7.driver import force_lhe_output_if_needed, run_generation +from madgraph.iolibs.template_files.mg7.model_widths import compute_auto_widths, load_process_model +from madgraph.iolibs.template_files.mg7.options import load_mg5_options +from madgraph.iolibs.template_files.mg7.postprocess import _off, _setup_logging + +class MG7Cmd(Cmd): + """The mg7 run interface: the command object that drives a launch. + + It plays two roles, and they are the same object on purpose. + + 1. It is the *mother* of the merged switch/card question (the thing + ``AskRunEditCard`` calls back into for ``keep_cards``/``do_open``/ + ``compute_widths``). + 2. It is the child command interface MG5 registers via + ``define_child_cmd_interface`` when you type ``launch`` on an mg7 output. + + Role 2 is what makes role 1 behave. ``Cmd.ask`` resolves a scripted answer + through ``self.check_answer_in_input_file``, where ``self`` is the mother -- + so once MG5 hands us its ``inputfile`` and registers itself as our + ``mother``, the question reads the very lines that follow ``launch`` in the + user's command file, and any line it does not recognise goes back up to MG5 + through ``store_line`` instead of being silently consumed. This is exactly + how a madevent output already behaves; before this class existed the mg7 + branch shelled out to ``bin/generate_events`` and had to guess which + following lines belonged to the run, which quietly swallowed every ``set``. + """ + + prompt = 'MG7> ' + + def __init__(self, me_dir='.', options=None, *args, **opts): + super().__init__(*args, **opts) + self.me_dir = os.path.abspath(me_dir) if me_dir != '.' else '.' + # The rest of the launcher is written against the process directory as + # the working directory, so do_generate_events chdirs there; but at + # construction time MG5 has not moved, hence the explicit me_dir. + self.options = load_mg5_options( + self.me_dir if self.me_dir != '.' else None) + if options: + # MG5's own resolved tool paths win: they are the ones the user + # configured in the session that is doing the launching. + self.options.update({k: v for k, v in options.items() + if k in self.options and v is not None}) + self.plugin_path = [] + self.proc_characteristics = {'grouped_matrix': False, 'limitations': []} + self._model = None # lazily imported by get_model() + + # ------------------------------------------------------------------ + # the contract the card question expects from its mother + # ------------------------------------------------------------------ + def keep_cards(self, need_card=[], ignore=[]): + from madgraph.interface.common_run_interface import CommonRunCmd + return CommonRunCmd.keep_cards(self, need_card, ignore) + + def do_open(self, line): + from madgraph.interface.common_run_interface import CommonRunCmd + CommonRunCmd.do_open(self, line) + + def check_open(self, args): + from madgraph.interface.common_run_interface import CommonRunCmd + CommonRunCmd.check_open(self, args) + + def get_model(self): + """The model this process was generated with (an ``import_ufo`` model). + + This is the hook the card question goes through for everything that + needs the model rather than the card: 'update dependent' (which resets + the masses/widths the model computes from the free parameters) and the + generic auto-width handler both call ``mother_interface.get_model()``. + Without it the launch question could only warn that it had failed to + update the dependent parameters, and MadSpin/the shower then read + whatever inconsistent values the param_card still held. + + None when the process recorded no model (an output made before + SubProcesses/model.txt existed) -- the caller warns in that case. + """ + if self._model is None: + me_dir = self.me_dir if self.me_dir != '.' else None + # False (not None) so a model that cannot be imported is not + # retried at every question. + self._model = load_process_model(me_dir) or False + return self._model or None + + def do_compute_widths(self, line): + # The interactive card editor delegates 'auto' width computation to + # the mother interface. Reuse the runtime helper (madgraph subprocess + # + the model stored at output time). ``line`` looks like + # " --path= [--nlo]"; we only need the card path. + m = re.search(r'--path=(\S+)', line or "") + path = m.group(1) if m else os.path.join("Cards", "param_card.dat") + compute_auto_widths(path) + # return an empty mapping: the caller iterates out.items() for the + # small-width treatment, which mg7 does not apply. + return {} + + # ------------------------------------------------------------------ + # running + # ------------------------------------------------------------------ + def do_generate_events(self, line): + """generate_events [-f] [-n NAME] [--laststep=parton] + + Run the mg7 generation and the selected post-processing tools.""" + + opts = self._parse_run_options(line) + cwd = os.getcwd() + target = self.me_dir if self.me_dir != '.' else cwd + try: + os.chdir(target) + _setup_logging() + if opts['name']: + self._set_run_name(opts['name']) + switch = {} + if not opts['force']: + switch = ask_edit_cards(mother=self) + switch = self._apply_laststep(switch, opts) + force_lhe_output_if_needed(switch) + _raise_open_file_limit() + run_generation(switch) + finally: + os.chdir(cwd) + + # ``launch`` is the name MG5 users type; keep it working here too. + do_launch = do_generate_events + + def do_npy_to_lhe(self, line): + """npy_to_lhe [RUN_NAME] [-o OUTPUT] [--seed N] [--no-gzip] + + Write the LHE file (events.lhe.gz) of a run that produced npy events, + e.g. to shower it. Without a run name, the most recent such run.""" + import shlex + me_dir = self.me_dir if self.me_dir != '.' else os.getcwd() + try: + npy_to_lhe.main(shlex.split(line), me_dir=me_dir) + except SystemExit as error: + if error.code not in (None, 0): + logger.error(str(error.code)) + + def do_quit(self, line): + """Leave the mg7 run interface.""" + return super().do_quit(line) + + do_exit = do_quit + + # ------------------------------------------------------------------ + # helpers + # ------------------------------------------------------------------ + def _parse_run_options(self, line): + """Understand the subset of MG5's `launch` options that mean something + for an mg7 run, and say so out loud for the ones that do not.""" + out = {'force': False, 'name': '', 'laststep': ''} + args = self.split_arg(line) + index = 0 + while index < len(args): + arg = args[index] + index += 1 + if arg in ('-f', '--force'): + out['force'] = True + elif arg in ('-n', '--name') and index < len(args): + out['name'] = args[index] + index += 1 + elif arg.startswith(('-n=', '--name=')): + out['name'] = arg.split('=', 1)[1] + elif arg in ('-s', '--laststep') and index < len(args): + out['laststep'] = args[index] + index += 1 + elif arg.startswith(('-s=', '--laststep=')): + out['laststep'] = arg.split('=', 1)[1] + elif arg.startswith('-'): + logger.warning( + "'%s' is not supported by the mg7 output and is ignored; " + "the equivalent settings live in Cards/run_card.toml", arg) + elif arg and not out['name']: + out['name'] = arg + return out + + def _set_run_name(self, name): + """Honour `launch -n NAME`: the mg7 run directory is + Events/_NN, so the name is a run_card setting.""" + path = os.path.join("Cards", "run_card.toml") + run_card = RunCardMG7(path, consistency=False) + if run_card["run"]["run_name"] != name: + run_card["run"]["run_name"] = name + run_card.write(path) + logger.info("run name set to '%s'", name) + + def _apply_laststep(self, switch, opts): + """`--laststep=parton` means: generate the events, run nothing on them. + + The mg7 tool selection is the switch dict the question returns, so + stopping at parton level is switching those entries off. + """ + if opts['laststep'] not in ('parton', 'auto', ''): + logger.warning( + "--laststep=%s is not supported by the mg7 output; " + "select the programs to run in the launch question instead", + opts['laststep']) + return switch + if opts['laststep'] != 'parton' or not switch: + return switch + switch = dict(switch) + for key in ('shower', 'detector', 'analysis', 'madspin', 'reweight'): + if not _off(switch.get(key, 'OFF')): + logger.info("--laststep=parton: not running %s", key) + switch[key] = 'OFF' + return switch + + +def build_selector_cmd(mother=None): + """Build the (monkey-patched) mother command + merged switch/card selector + used by the mg7 output. Returns the selector *class* and a mother instance + understood by AskRun/AskforEditCard. + + ``mother`` lets an in-process caller supply its own :class:`MG7Cmd` -- the + one MG5 registered as its child -- so that the question resolves scripted + answers against MG5's own command file (see :class:`MG7Cmd`). When it is + omitted a standalone instance is built, which is what ``bin/generate_events`` + gets. + """ + + from madgraph.interface.madevent_interface import AskRunEditCard + + from madgraph.various import banner as _banner_mod + from madgraph.various import misc as _misc + + class MG7Selector(AskRunEditCard): + """Merged switch/card question for the mg7 output. + + The mg7 run_card is a TOML file, so define_paths/init_run/do_set are + overridden *on this class* (not globally on AskforEditCard) to treat it + as such. Scoping them here is important: MadSpin/reweight spawn their own + legacy madevent runs in the same process, and a global patch would break + their (run_card.dat based) card editing.""" + + # param_card + the TOML run_card are always offered + always_cards = ['param_card.dat', 'run_card.toml'] + optional_cards = [] + + def define_paths(self, **opt): + super().define_paths(**opt) + self.paths["run"] = os.path.join(self.me_dir, "Cards", "run_card.toml") + self.paths["run_card.toml"] = os.path.join(self.me_dir, "Cards", "run_card.toml") + # the TOML run_card uses its own default file (concrete defaults + # written at output time); this powers "set default". + self.paths["run_default"] = os.path.join(self.me_dir, "Cards", "run_card_default.toml") + + def init_run(self, cards): + # Make sure the run_card is loaded as a RunCardMG7 (the generic + # editor does not recognise run_card.toml), else "set " fails. + out = super().init_run(cards) + if not isinstance(getattr(self, "run_card", None), RunCardMG7): + toml_path = self.paths.get("run") or os.path.join( + self.me_dir, "Cards", "run_card.toml") + if os.path.exists(toml_path): + try: + # allow_scan so a run_card holding scan:[...] values loads + with _misc.TMP_variable(_banner_mod.RunCard, "allow_scan", True): + self.run_card = RunCardMG7(toml_path, consistency="warning") + self.run_set = list(self.run_card.keys()) + except Exception as err: + logger.warning("could not load %s: %s", toml_path, err) + if isinstance(getattr(self, "run_card", None), RunCardMG7): + self.run_card.allow_scan = True + return getattr(self, "run_set", out) + + def do_set(self, line, *args, **kwargs): + # madevent-style shortcuts (lhc/lep/fixed_scale/no_parton_cut), cut + # editing, energy units and arithmetic/mass expressions on the TOML + # run_card, intercepted before delegating to the generic editor. + targs = self.split_arg(line) + run_card = getattr(self, "run_card", None) + if isinstance(run_card, RunCardMG7) and targs: + start = 1 if targs[0] == "run_card" else 0 + if len(targs) > start: + name = targs[start] + nlow = name.lower() + rest = " ".join(targs[start + 1:]).split("#")[0].strip() + masses = run_card.get_mass_shortcuts(getattr(self, "param_card", None)) + + if nlow in ("no_parton_cut", "nocut", "no_cut"): + run_card.remove_all_cut() + logger.info("removing all cuts from the run_card.toml") + self.modified_card.add("run") + return + # [histograms] is a list of observables, not a parameter, + # so "set histograms ..." can only mean the whole section. + if nlow in ("histograms", "histogram", "plots"): + value = rest.lower() + if value in ("off", "none", "no", "false", "0"): + run_card.remove_all_histograms() + logger.info( + "removing all histograms from the run_card.toml") + self.modified_card.add("run") + return + if value in ("default", "on"): + if run_card.restore_default_histograms( + self.paths.get("run_default")): + logger.info( + "restored the %d histograms written at " + "output time", + len(run_card["histograms"])) + self.modified_card.add("run") + else: + logger.warning( + "no run_card_default.toml to restore the " + "histograms from") + return + logger.warning( + "'set histograms %s': only OFF (remove them all) " + "and default (put the generated ones back) are " + "understood; edit the [histograms] section of " + "Cards/run_card.toml for anything else", + rest or "(no value)") + return + if nlow in ("lhc", "lep", "ilc", "lcc") and rest: + ecm = run_card.set_collider(nlow, rest, masses) + logger.info("set %s collider: e_cm = %s GeV", nlow, ecm) + self.modified_card.add("run") + return + if nlow == "fixed_scale" and rest: + val = run_card.set_fixed_scale(rest, masses) + logger.info("set fixed scales to %s GeV", val) + self.modified_card.add("run") + return + + if rest and run_card.is_cut_name(name): + cut, bound, val = run_card.set_cut(name, run_card.evaluate(rest, masses)) + logger.info("modify cut %s.%s of the run_card.toml to %s", cut, bound, val) + self.modified_card.add("run") + return + + # 'iseed' is the madevent spelling of the mg7 run.seed, + # but the two disagree on how to ask for a random seed: + # madevent uses 0, mg7 uses -1 (0 being a perfectly good + # fixed seed there). Translate here rather than in + # _LO_SCALAR_MAP, which also drives the LO -> MG7 run_card + # conversion and would silently pin those runs to seed 0. + if rest and nlow == 'iseed': + value = run_card.evaluate(rest, masses) + try: + value = int(value) + except (TypeError, ValueError): + value = None + if value is None: + logger.warning("ignoring 'set iseed %s': not an integer", rest) + return + seed = -1 if value == 0 else value + run_card.set('run.seed', seed, user=True) + logger.info("set iseed (mg7 run.seed) of the run_card.toml to %s%s", + seed, " (random)" if seed == -1 else "") + self.modified_card.add("run") + return + + # legacy madevent run_card names -> mg7 "section.key", so a + # madevent-style launch script ("set nevents 500", "set + # use_syst F", "set bwcutoff 10", ...) edits the mg7 + # run_card.toml verbatim. RunCardMG7._LO_SCALAR_MAP is the + # same rename table used by the LO->MG7 run_card conversion. + if rest and nlow in run_card._LO_SCALAR_MAP \ + and nlow not in [k.lower() for k in run_card.keys()]: + target = run_card._LO_SCALAR_MAP[nlow] + run_card.set(target, rest, user=True) + logger.info("set %s (mg7 %s) of the run_card.toml to %s", + nlow, target, run_card[target]) + self.modified_card.add("run") + return + + if rest and nlow in [k.lower() for k in run_card.keys()]: + current = run_card[nlow] + if isinstance(current, (int, float)) and not isinstance(current, bool): + resolved = run_card.evaluate(rest, masses) + if not isinstance(resolved, str): + prefix = "run_card " if start == 1 else "" + line = "%s%s %s" % (prefix, name, resolved) + return super().do_set(line, *args, **kwargs) + + return MG7Selector, (mother if mother is not None else MG7Cmd()) + + +def ask_edit_cards(mother=None) -> dict: + """Single (MadDM-style) question letting the user both pick which programs + to run after generation and edit the associated cards. Returns the switch + dict describing the selected tools. + + ``mother`` is the :class:`MG7Cmd` asking the question. Passing the instance + MG5 registered as its child is what lets a scripted ``launch`` answer this + question from MG5's own command file. + """ + + selector_class, mother = build_selector_cmd(mother) + # path_msg is what makes Cmd.check_answer_in_input_file accept a bare path as + # an answer (its "elif path:" branch), so that a scripted launch can hand the + # question a card/banner path -- as the question itself advertises -- and have + # it replace the corresponding card. Without it the path is rejected ("This + # answer is not valid for current question") and the default is used instead. + try: + switch, question = mother.ask('', '0', [], path_msg='enter path', + ask_class=selector_class, + mode='auto', line_args=[], force=False, + return_instance=True) + except BaseException: + # interrupted (e.g. ctrl-C) before the prune below could run: prune + # here too, using the instance ask() stashed pre-interrupt, or the + # eagerly-materialised tool cards get misread as "left ON" next time. + instance = getattr(mother, '_last_ask_instance', None) + if instance is not None and hasattr(instance, 'active_cards'): + try: + mother.keep_cards(instance.active_cards()) + except Exception as error: + logger.debug('could not prune tool cards after aborted launch question: %s', error) + raise + switch = dict(switch) + prune_unselected_tool_cards(mother, question, switch) + return switch + + +def prune_unselected_tool_cards(mother, question, switch) -> None: + """Hide the cards of the tools that were NOT selected in the question. + + The question materialises every candidate tool card (copying the + *_default.dat) so that it can offer them all for edition; madevent's + ask_run_configuration prunes them again afterwards (keep_cards on the + selected cards), but the mg7 launcher had no such cleanup. The materialised + cards therefore persisted, and a later re-launch of the same output saw them + all present and defaulted every switch to ON (set_default_ keys off + card presence). Prune here too, using the question's own switch->card map, so + a re-launch reflects the previous selection instead of turning everything on. + """ + try: + keep = list(question.always_cards) + for spec in question.switch_cards: + if spec['on'](switch): + keep.append(spec['card']) + mother.keep_cards(keep) + except Exception as error: + logger.debug('could not prune unselected tool cards: %s', error) + + +def _raise_open_file_limit() -> None: + """Remove the soft limit on the number of open files: it can be quite low + on some systems and the event-combination step opens one file per channel.""" + soft_lim, hard_lim = resource.getrlimit(resource.RLIMIT_NOFILE) + try: + resource.setrlimit(resource.RLIMIT_NOFILE, (hard_lim, hard_lim)) + except (ValueError, OSError) as error: + logger.debug("could not raise the open-file limit: %s", error) + + +def main() -> None: + """``bin/generate_events`` entry point. + + It runs the same :class:`MG7Cmd` command MG5's in-process ``launch`` uses, + so the two entry points cannot drift apart; the only difference is that + nothing here supplies a ``mother``/``inputfile``, so the question falls back + to reading this process's own stdin. + """ + parser = argparse.ArgumentParser() + parser.add_argument("-f", action="store_false", dest="ask_edit_cards") + args = parser.parse_args() + + MG7Cmd().run_cmd("generate_events" if args.ask_edit_cards else "generate_events -f") diff --git a/madgraph/iolibs/template_files/mg7/common.py b/madgraph/iolibs/template_files/mg7/common.py new file mode 100644 index 0000000000..253f4d06d9 --- /dev/null +++ b/madgraph/iolibs/template_files/mg7/common.py @@ -0,0 +1,200 @@ +"""Shared state of the mg7 launcher: the madspace import, logging and small helpers.""" + +import logging +import os +import re +import subprocess +import sys +import time + +# Install madspace on first use, then import it. The bootstrap lives in its own +# module (which must not import madspace) so that an in-process caller can run +# it *before* importing this one -- see bootstrap.ensure_madspace. +from madgraph.iolibs.template_files.mg7.bootstrap import ( + ensure_madspace as _ensure_madspace, + drop_install_path as _drop_install_path, + MADSPACE_DIR as _MADSPACE_DIR, +) + +_ensure_madspace() + +import madspace as ms + +# madspace is imported; stop its install directory from shadowing the caller's +# yaml/packaging/... for the rest of what is now MG5's own session. +_drop_install_path() + +_source_hash = subprocess.run( + [sys.executable, str(_MADSPACE_DIR / "source_hash.py")], + capture_output=True, text=True, check=True, +).stdout.strip() +if _source_hash != ms.SOURCE_HASH: + print() + print( + "\033[1m\033[31mWARNING\033[39m: madspace source and installed binaries " + "are not compatible (source hash mismatch) — consider recompiling " + "madspace (e.g. `install madspace -y`)\033[0m" + ) + print() + +logger = logging.getLogger("madgraph7") +LOG_LEVEL_MAP = { + ms.Logger.LogLevel.level_debug: logging.DEBUG, + ms.Logger.LogLevel.level_info: logging.INFO, + ms.Logger.LogLevel.level_warning: logging.WARNING, + ms.Logger.LogLevel.level_error: logging.ERROR, +} +def ms_log_handler(level: ms.Logger.LogLevel, message: str): + logger.log(LOG_LEVEL_MAP[level], message) +ms.Logger.set_log_handler(ms_log_handler) + + +def get_start_time(): + return time.time(), time.process_time() + + +def format_time(t: int, centi: bool = False): + hours, t = divmod(t, 3600) + minutes, seconds = divmod(t, 60) + if centi: + return f"{int(hours):02}:{int(minutes):02}:{seconds:02.2f}" + else: + return f"{int(hours):02}:{int(minutes):02}:{seconds:02.0f}" + + +def resolve_verbosity(verbosity: str) -> str: + """Resolve the run_card "auto" verbosity to "pretty"/"log" depending on + whether stdout is attached to a terminal; other values pass through + unchanged.""" + if verbosity == "auto": + return "pretty" if sys.stdout.isatty() else "log" + return verbosity + + +def resolve_seed(seed: int) -> int: + """Resolve the run_card "seed": -1 draws a fresh 64-bit seed via + os.urandom, any other value is used as-is.""" + if seed == -1: + return int.from_bytes(os.urandom(8), "big") + return seed + + +def device_type_of(device_name: str) -> str: + """The device type of one 'device' run_card entry (the optional ':' + suffix is stripped).""" + return device_name.split(":")[0] + + +def backend_of(device_name: str, cpu_mode: str) -> str: + """The build BACKEND that serves one 'device' run_card entry. + + ``cpu_mode`` describes the SIMD width of the CPU code and is therefore + meaningless for the cuda/hip devices: those build the backend named after + the device itself. + """ + device_type = device_type_of(device_name) + return cpu_mode if device_type == "cpu" else device_type + + +_PRECISION_2_FPTYPE = {"all64": "d", "all32": "f", "color32": "m", "denom64": "v"} + + +def fptype_of(precision: str) -> str: + return _PRECISION_2_FPTYPE.get(str(precision).lower(), str(precision).lower()) + + +def resolve_auto_backend(build_path: str) -> str: + """Ask the matrix-element Makefile to resolve the ``auto`` cpu_mode. + + Given the produced shared library will have its name taken from the resolved + backend name, we need to make sure the detection is taking place correctly + and catch any possible error. + """ + command = ["make", "-n", "BACKEND=auto", "detect-backend"] + try: + result = subprocess.run( + command, + cwd=build_path, + capture_output=True, + text=True, + check=True, + ) + except OSError as exc: + raise RuntimeError( + f"Could not run make to resolve cpu_mode='auto' in '{build_path}': {exc}" + ) from exc + except subprocess.CalledProcessError as exc: + output = "\n".join( + part.strip() for part in (exc.stdout, exc.stderr) if part.strip() + ) + detail = f"\nmake output:\n{output}" if output else "" + raise RuntimeError( + f"Could not resolve cpu_mode='auto' in '{build_path}': " + f"Exit status {exc.returncode}.{detail}" + ) from exc + + match = re.search( + r"^BACKEND=(\S+) \(was auto\)$", result.stdout, re.MULTILINE + ) + if match is None or match.group(1) == "auto": + output = "\n".join( + part.strip() for part in (result.stdout, result.stderr) if part.strip() + ) + detail = f"\nmake output:\n{output}" if output else "" + raise RuntimeError( + f"Could not resolve cpu_mode='auto' in '{build_path}': " + f"make failed to report a backend.{detail}" + ) + return match.group(1) + + +def me_parameters(run_card): + """Run-time parameters of the matrix-element libraries, passed to each + instance through umami_set_parameter by ms.Context.load_matrix_element: + the window of the $-excluded propagators is the run card bw_cutoff (the + one the phase space uses around the resonances too), and with + [generation] interference_helicity = "summed" the helicity sum as the + |M|^2 of an interference (see interference_helicities_summed). That one is + only sent when asked for: a library made before it existed refuses an + unknown parameter.""" + parameters = {"bwcutoff": float(run_card["phasespace"]["bw_cutoff"])} + if interference_helicities_summed(run_card): + parameters["interference_helicity_summed"] = 1. + return parameters + + +def interference_helicities_summed(run_card): + """Whether [generation] interference_helicity is "summed". + + For an interference (squared split orders dropping a component), a helicity + can contribute negatively. "exact", the default, draws the helicity of each + event on |T_i| and gives the event sign(T_i) * sum_j |T_j|, as madevent does: + the helicity of the event is exact, at the cost of a larger sum|w| (about 2x + the events for the same precision). "summed" keeps the helicity sum sum_j T_j + as the weight (lower variance), and the helicity then means nothing: it is + written as 9 in the LHE (lhe_helicities).""" + return run_card["generation"].get("interference_helicity", "exact") == "summed" + + +def lhe_helicities(meta, run_card): + """The helicity table the LHE completer writes for a subprocess: its own, + or 9 for every particle of an interference subprocess whose |M|^2 is the + helicity sum (interference_helicities_summed).""" + helicities = meta["helicities"] + if meta.get("interference") and interference_helicities_summed(run_card): + return [[9] * len(row) for row in helicities] + return helicities + + +def lhe_weight_info(status): + """(IDWTUP, XMAXUP) for the LHE block of a run with this generator + status. The events are unweighted to +-sigma_abs, keeping the sign of the + integrand, and XMAXUP is that unit weight rather than the signed cross + section. A sample that can hold negative weights -- an interference -- + declares it with IDWTUP = -4, as madevent does: the cross section is then + the mean of the signed event weights. With -3 a shower takes it as + |XSECUP| times the mean sign instead (Pythia8 does), which is not the + cross section, and +3 is for positive weights only. gridpack.py applies the + same rule inline, since it must not import this module.""" + negative = status.mean < 0 or status.mean_abs > abs(status.mean) * (1 + 1e-12) + return (-4 if negative else 3), status.mean_abs diff --git a/madgraph/iolibs/template_files/mg7/driver.py b/madgraph/iolibs/template_files/mg7/driver.py new file mode 100644 index 0000000000..ed1d974ce9 --- /dev/null +++ b/madgraph/iolibs/template_files/mg7/driver.py @@ -0,0 +1,186 @@ +"""Runs, scans and the entry point of the mg7 launcher.""" + +import gc +import logging +import os +import re +import shutil + +from madgraph.iolibs.template_files.mg7.common import logger +from madgraph.various.banner import RunCardMG7 +from madgraph.iolibs.template_files.mg7.model_widths import compute_auto_widths +from madgraph.iolibs.template_files.mg7.postprocess import run_lhe_postprocessing, run_selected_tools, selected_tools +from madgraph.iolibs.template_files.mg7.madgraph_process import MadgraphProcess + +def _release_channel_generators(process) -> None: + """Drop the per-channel event/weight generators so their intermediate + .npy files are deleted right away. madspace only frees them when the + owning Python objects are collected, and process.event_generator / + process.phasespaces sit in a reference cycle that plain refcounting + never breaks -- only an explicit gc.collect() does, promptly, instead of + leaving it to whenever the cyclic GC next runs on its own.""" + process.event_generator = None + process.phasespaces = None + gc.collect() + + +def run_single(switch=None) -> "MadgraphProcess": + """Run a single generation and return the process (for its result).""" + compute_auto_widths() + process = MadgraphProcess() + process.survey() + process.train_madnis() + process.generate_events() + if process.gridpack_run_mode != "regular": + return process + # run_card-driven LHE post-processing (displaced vertex + systematics) + run_lhe_postprocessing(process) + # run the post-processing tools (Pythia8/Delphes/MadSpin/reweight/analysis) + # selected in the merged question above on the generated events. + if switch: + run_selected_tools(switch, process) + return process + + +def detect_run_scan(run_card_path): + """Return a banner.RunCardIterator if the run_card contains scan:[...] + values, else None.""" + from madgraph.various import banner as banner_mod + from madgraph.various import misc as _misc + with _misc.TMP_variable(banner_mod.RunCard, 'allow_scan', True): + rc = banner_mod.RunCard(run_card_path, consistency=False) + if getattr(rc, 'scan_set', None): + return banner_mod.RunCardIterator(run_card_path) + return None + + +def detect_param_scan(param_card_path): + """Return a ParamCardIterator if the param_card contains scan:[...] values, + else None.""" + if not os.path.exists(param_card_path): + return None + from models import check_param_card as param_card_mod + it = param_card_mod.ParamCardIterator(param_card_path) + for block in it.order: + for param in block: + if isinstance(param.value, str) and param.value.strip().lower().startswith('scan'): + return it + return None + + +def run_scan(iterator, card_path, switch=None) -> None: + """Iterate over all scan points, running a full generation for each and + accumulating the results, then write the scan summary. Works for both the + run_card (RunCardIterator) and the param_card (ParamCardIterator); their + interface (__iter__/write/store_entry/get_next_name/write_summary) is the + same. The scan card is restored afterwards.""" + import tomllib + with open(os.path.join("Cards", "run_card.toml"), "rb") as f: + run_name = tomllib.load(f).get("run", {}).get("run_name", "run") + + from models import check_param_card as param_card_mod + is_param_scan = isinstance(iterator, param_card_mod.ParamCardIterator) + + backup = card_path + ".scan_bak" + shutil.copy(card_path, backup) + try: + for i, point in enumerate(iterator): + point.write(card_path) + logger.info("=== scan point %d ===", i + 1) + process = run_single(switch) + # use the run directory the process actually created, so the + # per-point params.dat written by write_summary has a home + name = os.path.basename(process.run_path) + if is_param_scan: + # pass the (possibly auto-width-updated) card so the summary + # records the model-computed width for each 'auto' entry + iterator.store_entry(name, process.get_result(), + param_card_path=card_path) + else: + iterator.store_entry(name, process.get_result()) + _release_channel_generators(process) + os.makedirs("Events", exist_ok=True) + summary = os.path.join("Events", "scan_%s.txt" % run_name) + iterator.write_summary(summary) + logger.info("scan results written to %s", summary) + finally: + shutil.move(backup, card_path) + + +def run_generation(switch=None) -> None: + """Run the generation, expanding a scan over the run_card or the param_card + when one is present (scanning both simultaneously is not allowed).""" + run_card_path = os.path.join("Cards", "run_card.toml") + param_card_path = os.path.join("Cards", "param_card.dat") + run_iter = detect_run_scan(run_card_path) + param_iter = detect_param_scan(param_card_path) + if run_iter and param_iter: + raise RuntimeError( + "Scanning simultaneously over the run_card and the param_card is " + "not allowed. Please keep the scan:[...] entries in only one card.") + if run_iter: + run_scan(run_iter, run_card_path, switch) + elif param_iter: + run_scan(param_iter, param_card_path, switch) + else: + _release_channel_generators(run_single(switch)) + + +def _lhe_needed_by(switch, card) -> "str | None": + """Why the run needs an LHE event file, or None when it does not: every + post-processing tool run_selected_tools drives and the run_card + [postprocessing] steps read the LHE file. ``card`` is the raw TOML content + of the run_card.""" + tools = selected_tools(switch) if switch else [] + if tools: + return tools[0] + cfg = card.get("postprocessing", {}) + try: + if float(cfg.get("time_of_flight", -1.0)) >= 0: + return "time_of_flight" + except (TypeError, ValueError): + pass + # the legacy systematics.py path only runs when madspace did not compute + # the weights itself (see run_lhe_postprocessing) + if cfg.get("systematics") \ + and not card.get("systematics", {}).get("enable", True): + return "systematics" + return None + + +def force_lhe_output_if_needed(switch) -> None: + """The events are written as npy by default; switch the run_card to the + LHE format when a selected post-processing needs an LHE file. The card is + rewritten rather than refused, so that asking for a shower in the launch + question is enough -- the user does not have to know that it constrains + output_format. + + The card is read as raw TOML and only its output_format line is edited: a + run_card scan ("scan:[...]" values) is not a valid RunCardMG7 and must + survive untouched for run_scan.""" + import tomllib + path = os.path.join("Cards", "run_card.toml") + if not os.path.exists(path): + return + with open(path, "rb") as f: + card = tomllib.load(f) + output_format = card.get("run", {}).get( + "output_format", RunCardMG7()["run"]["output_format"]) + if output_format == "lhe": + return + reason = _lhe_needed_by(switch, card) + if reason is None: + return + with open(path) as f: + text = f.read() + text, count = re.subn(r'(?m)^(\s*output_format\s*=\s*)("[^"]*"|\'[^\']*\'|\S+)', + r'\g<1>"lhe"', text, count=1) + if not count: + text, count = re.subn(r'(?m)^\[run\][^\n]*\n', '\\g<0>output_format = "lhe"\n', + text, count=1) + if not count: + text += '\n[run]\noutput_format = "lhe"\n' + with open(path, "w") as f: + f.write(text) + logging.getLogger("madevent").info( + "output_format set to 'lhe' (required by %s).", reason) diff --git a/madgraph/iolibs/template_files/mg7/launch.py b/madgraph/iolibs/template_files/mg7/launch.py index b871e59039..ceb4420a49 100644 --- a/madgraph/iolibs/template_files/mg7/launch.py +++ b/madgraph/iolibs/template_files/mg7/launch.py @@ -1,4685 +1,73 @@ -import argparse -import gc -import os -import sys -import time -from datetime import timedelta -from pathlib import Path -import glob -import shutil -import json -import subprocess -import re -import logging -from dataclasses import dataclass, field -from typing import Literal, NamedTuple -import resource - -# Install madspace on first use, then import it. The bootstrap lives in its own -# module (which must not import madspace) so that an in-process caller can run -# it *before* importing this one -- see bootstrap.ensure_madspace. -from madgraph.iolibs.template_files.mg7.bootstrap import ( - ensure_madspace as _ensure_madspace, - drop_install_path as _drop_install_path, - MG_ROOT as _MG_ROOT, - MADSPACE_DIR as _MADSPACE_DIR, - INSTALL_DIR as _INSTALL_DIR, +"""Entry point of the mg7 launcher, re-exporting its modules.""" + +from madgraph.iolibs.template_files.mg7.common import * # noqa: F401,F403 +from madgraph.iolibs.template_files.mg7.common import ms, logger # noqa: F401 +from madgraph.iolibs.template_files.mg7.observables import ( # noqa: F401 + CutItem, + HistItem, + _MERGED_PID_REPRESENTATIVE, + clean_pids, + _SIGNED_OBSERVABLES, + _flavor_options, + observable_pids, + build_observable, + decay_products, + pid_is_qcd, ) - -_ensure_madspace() - -import madspace as ms - -# madspace is imported; stop its install directory from shadowing the caller's -# yaml/packaging/... for the rest of what is now MG5's own session. -_drop_install_path() -from models.check_param_card import ParamCard -from madgraph.iolibs.template_files.mg7 import hwu_output -from madgraph.iolibs.template_files.mg7 import plots -from madgraph.iolibs.template_files.mg7 import systematics_summary -from madgraph.iolibs.template_files.mg7 import npy_to_lhe -from madgraph.iolibs.template_files.mg7.custom_function_loader import ( - load_custom_function, - load_surrogate, +from madgraph.iolibs.template_files.mg7.subprocess_madmatrix import ( # noqa: F401 + Channel, + PhaseSpace, + MultiChannelData, + build_topologies, + build_multi_channel_data, + MadgraphSubprocess, +) +from madgraph.iolibs.template_files.mg7.madgraph_process import ( # noqa: F401 + MadgraphProcess, +) +from madgraph.iolibs.template_files.mg7.options import ( # noqa: F401 + _ROOTED_OPTIONS, + load_mg5_options, + lhapdf_paths, +) +from madgraph.iolibs.template_files.mg7.model_widths import ( # noqa: F401 + _model_hash_warned, + _model_cache, + read_stored_model, + _proc_characteristic, + load_process_model, + compute_auto_widths, +) +from madgraph.iolibs.template_files.mg7.postprocess import ( # noqa: F401 + _find_event_file, + _find_or_convert_event_file, + _report_failure, + _off, + _has_effective_handler, + _TOOL_LOGGING_READY, + _setup_logging, + selected_tools, + run_selected_tools, + _add_time_of_flight, + _lhapdf_config_path, + _run_systematics, + run_lhe_postprocessing, +) +from madgraph.iolibs.template_files.mg7.driver import ( # noqa: F401 + _release_channel_generators, + run_single, + detect_run_scan, + detect_param_scan, + run_scan, + run_generation, + _lhe_needed_by, + force_lhe_output_if_needed, +) +from madgraph.iolibs.template_files.mg7.cli import ( # noqa: F401 + MG7Cmd, + build_selector_cmd, + ask_edit_cards, + prune_unselected_tool_cards, + _raise_open_file_limit, + main, ) -from madgraph.various.banner import RunCardMG7 -from madgraph.various import misc -from madgraph.interface.extended_cmd import Cmd - -_source_hash = subprocess.run( - [sys.executable, str(_MADSPACE_DIR / "source_hash.py")], - capture_output=True, text=True, check=True, -).stdout.strip() -if _source_hash != ms.SOURCE_HASH: - print() - print( - "\033[1m\033[31mWARNING\033[39m: madspace source and installed binaries " - "are not compatible (source hash mismatch) — consider recompiling " - "madspace (e.g. `install madspace -y`)\033[0m" - ) - print() - -logger = logging.getLogger("madgraph7") -LOG_LEVEL_MAP = { - ms.Logger.LogLevel.level_debug: logging.DEBUG, - ms.Logger.LogLevel.level_info: logging.INFO, - ms.Logger.LogLevel.level_warning: logging.WARNING, - ms.Logger.LogLevel.level_error: logging.ERROR, -} -def ms_log_handler(level: ms.Logger.LogLevel, message: str): - logger.log(LOG_LEVEL_MAP[level], message) -ms.Logger.set_log_handler(ms_log_handler) - - -def get_start_time(): - return time.time(), time.process_time() - - -def format_time(t: int, centi: bool = False): - hours, t = divmod(t, 3600) - minutes, seconds = divmod(t, 60) - if centi: - return f"{int(hours):02}:{int(minutes):02}:{seconds:02.2f}" - else: - return f"{int(hours):02}:{int(minutes):02}:{seconds:02.0f}" - - -def resolve_verbosity(verbosity: str) -> str: - """Resolve the run_card "auto" verbosity to "pretty"/"log" depending on - whether stdout is attached to a terminal; other values pass through - unchanged.""" - if verbosity == "auto": - return "pretty" if sys.stdout.isatty() else "log" - return verbosity - - -def resolve_seed(seed: int) -> int: - """Resolve the run_card "seed": -1 draws a fresh 64-bit seed via - os.urandom, any other value is used as-is.""" - if seed == -1: - return int.from_bytes(os.urandom(8), "big") - return seed - - -def device_type_of(device_name: str) -> str: - """The device type of one 'device' run_card entry (the optional ':' - suffix is stripped).""" - return device_name.split(":")[0] - - -def backend_of(device_name: str, cpu_mode: str) -> str: - """The build BACKEND that serves one 'device' run_card entry. - - ``cpu_mode`` describes the SIMD width of the CPU code and is therefore - meaningless for the cuda/hip devices: those build the backend named after - the device itself. - """ - device_type = device_type_of(device_name) - return cpu_mode if device_type == "cpu" else device_type - - -_PRECISION_2_FPTYPE = {"all64": "d", "all32": "f", "color32": "m", "denom64": "v"} - - -def fptype_of(precision: str) -> str: - return _PRECISION_2_FPTYPE.get(str(precision).lower(), str(precision).lower()) - - -def resolve_auto_backend(build_path: str) -> str: - """Ask the matrix-element Makefile to resolve the ``auto`` cpu_mode. - - Given the produced shared library will have its name taken from the resolved - backend name, we need to make sure the detection is taking place correctly - and catch any possible error. - """ - command = ["make", "-n", "BACKEND=auto", "detect-backend"] - try: - result = subprocess.run( - command, - cwd=build_path, - capture_output=True, - text=True, - check=True, - ) - except OSError as exc: - raise RuntimeError( - f"Could not run make to resolve cpu_mode='auto' in '{build_path}': {exc}" - ) from exc - except subprocess.CalledProcessError as exc: - output = "\n".join( - part.strip() for part in (exc.stdout, exc.stderr) if part.strip() - ) - detail = f"\nmake output:\n{output}" if output else "" - raise RuntimeError( - f"Could not resolve cpu_mode='auto' in '{build_path}': " - f"Exit status {exc.returncode}.{detail}" - ) from exc - - match = re.search( - r"^BACKEND=(\S+) \(was auto\)$", result.stdout, re.MULTILINE - ) - if match is None or match.group(1) == "auto": - output = "\n".join( - part.strip() for part in (result.stdout, result.stderr) if part.strip() - ) - detail = f"\nmake output:\n{output}" if output else "" - raise RuntimeError( - f"Could not resolve cpu_mode='auto' in '{build_path}': " - f"make failed to report a backend.{detail}" - ) - return match.group(1) - - -def me_parameters(run_card): - """Run-time parameters of the matrix-element libraries, passed to each - instance through umami_set_parameter by ms.Context.load_matrix_element: - the window of the $-excluded propagators is the run card bw_cutoff (the - one the phase space uses around the resonances too), and with - [generation] interference_helicity = "summed" the helicity sum as the - |M|^2 of an interference (see interference_helicities_summed). That one is - only sent when asked for: a library made before it existed refuses an - unknown parameter.""" - parameters = {"bwcutoff": float(run_card["phasespace"]["bw_cutoff"])} - if interference_helicities_summed(run_card): - parameters["interference_helicity_summed"] = 1. - return parameters - - -def interference_helicities_summed(run_card): - """Whether [generation] interference_helicity is "summed". - - For an interference (squared split orders dropping a component), a helicity - can contribute negatively. "exact", the default, draws the helicity of each - event on |T_i| and gives the event sign(T_i) * sum_j |T_j|, as madevent does: - the helicity of the event is exact, at the cost of a larger sum|w| (about 2x - the events for the same precision). "summed" keeps the helicity sum sum_j T_j - as the weight (lower variance), and the helicity then means nothing: it is - written as 9 in the LHE (lhe_helicities).""" - return run_card["generation"].get("interference_helicity", "exact") == "summed" - - -def lhe_helicities(meta, run_card): - """The helicity table the LHE completer writes for a subprocess: its own, - or 9 for every particle of an interference subprocess whose |M|^2 is the - helicity sum (interference_helicities_summed).""" - helicities = meta["helicities"] - if meta.get("interference") and interference_helicities_summed(run_card): - return [[9] * len(row) for row in helicities] - return helicities - - -def lhe_weight_info(status): - """(IDWTUP, XMAXUP) for the LHE block of a run with this generator - status. The events are unweighted to +-sigma_abs, keeping the sign of the - integrand, and XMAXUP is that unit weight rather than the signed cross - section. A sample that can hold negative weights -- an interference -- - declares it with IDWTUP = -4, as madevent does: the cross section is then - the mean of the signed event weights. With -3 a shower takes it as - |XSECUP| times the mean sign instead (Pythia8 does), which is not the - cross section, and +3 is for positive weights only. gridpack.py applies the - same rule inline, since it must not import this module.""" - negative = status.mean < 0 or status.mean_abs > abs(status.mean) * (1 + 1e-12) - return (-4 if negative else 3), status.mean_abs - - -@dataclass -class Channel: - phasespace_mapping: ms.PhaseSpaceMapping - adaptive_mapping: ms.Flow | ms.VegasMapping - discrete_sym: ms.DiscreteSampler | ms.DiscreteFlow | None - discrete_flavor: ms.DiscreteSampler | ms.DiscreteFlow | None - channel_weight_indices: list[int] | None - name: str - active_flavors: list[int] - event_generator: ms.ChannelEventGenerator | None = None - - -@dataclass -class PhaseSpace: - mode: Literal["multichannel", "flat", "both"] - channels: list[Channel] - symfact: list[int | None] - first_chan_weight_remap: list[list[int]] = field(default_factory=list) - first_remapped_chan_count: int = 0 - second_chan_weight_remap: list[int] = field(default_factory=list) - second_remapped_chan_count: int = 0 - prop_chan_weights: ms.PropagatorChannelWeights | None = None - subchan_weights: ms.SubchannelWeights | None = None - cwnet: ms.ChannelWeightNetwork | None = None - - -class MultiChannelData(NamedTuple): - amp2_remaps: list[list[int]] - symfact: list[int | None] - topologies: list[list[ms.Topology]] - permutations: list[list[list[int]]] - channel_indices: list[list[int]] - channel_weight_indices: list[list[list[int]]] - diagram_indices: list[list[int]] - diagram_color_indices: list[list[list[int]]] - diagram_propagator_pdgs: list[list[list[int]]] - active_flavors: list[list[list[int]]] - qcd_s_channel_count: list[int] - - -@dataclass -class CutItem: - observable_kwargs: dict - min: float - max: float - mode: str - - -@dataclass -class HistItem: - observable_kwargs: dict - min: float - max: float - bin_count: int - # the distribution of the event weight itself (the reserved "weight" key - # of [histograms]); it is read off the event record, not computed from - # the momenta, so it has no observable - name: str = "" - from_weight: bool = False - - -class MadgraphProcess: - def __init__(self): - self.madnis_trained = False - self.lhe_completer = None - self.load_cards() - self.init_backend() - self.init_event_dir() - self.init_context() - self.init_custom_functions() - self.init_surrogate() - self.init_cuts() - self.init_histograms() - self.init_generator_config() - self.init_beam() - self.init_subprocesses() - - def load_cards(self) -> None: - self.run_card = RunCardMG7(os.path.join("Cards", "run_card.toml")) - # Resolved once so every generator built during this run shares the same - # seed; the concrete value (even if randomly drawn) is recorded in each - # generator's info.json. - self.run_seed = resolve_seed(self.run_card["run"]["seed"]) - self.param_card_path = os.path.join("Cards", "param_card.dat") - self.param_card = ParamCard(self.param_card_path) - with open(os.path.join("SubProcesses", "subprocesses.json")) as f: - self.subprocess_data = json.load(f) - if self.run_card["phasespace"]["merge_subprocesses"]: - with open(os.path.join("SubProcesses", "merged_subprocesses.json")) as f: - self.merged_subprocess_data = json.load(f) - else: - self.merged_subprocess_data = None - - self.init_decay_mode() - self.init_me_frame() - if interference_helicities_summed(self.run_card): - count = sum(1 for meta in self.subprocess_data if meta.get("interference")) - if count: - logger.info( - "interference_helicity = summed: the weight of the %d interference " - "subprocess(es) is the helicity sum, and their helicities are " - "written as 9 in the LHE", count) - - def init_me_frame(self) -> None: - """Resolve the run card's me_frame into the list of external particles - whose momentum sum defines the frame the matrix element is evaluated in. - - Empty (the default) means no boost: the matrix element sees the momenta - in the frame they are generated in. Unlike madevent, which hands its - matrix element momenta that are already in the partonic centre of mass, - madspace generates them in the lab frame, so asking for madevent's own - frame ([1, 2]) does mean an actual boost here. It changes nothing for a - Lorentz invariant matrix element (only the rounding), but it is what - defines the polarisation axes of a polarised one. - - The particle numbers are checked here rather than only in madspace, - which would not see them until the integrands are built -- long after - the matrix-element libraries have been compiled. Every subprocess is - checked against its own particles: an output can mix multiplicities - (p p > z j and p p > z j j, say, for MLM merging), and madspace builds - one matrix element per subprocess, each refusing a particle it does not - have. That is stricter than madevent, which ignores the particles a - process does not have and so changes the frame from one multiplicity - to the next without a word. - """ - me_frame = list(self.run_card["run"]["me_frame"]) - pdgs = [clean_pids(meta["incoming"]) + clean_pids(meta["outgoing"]) - for meta in self.subprocess_data] - counts = sorted({len(pdg_list) for pdg_list in pdgs}) - for meta, pdg_list in zip(self.subprocess_data, pdgs): - for index in me_frame: - if 1 <= index <= len(pdg_list): - continue - process = "%s > %s" % ( - " ".join(str(pid) for pid in meta["incoming"]), - " ".join(str(pid) for pid in meta["outgoing"]), - ) - mixed = ( - f"; this output mixes subprocesses with " - f"{', '.join(str(count) for count in counts)} external " - "particles, and me_frame has to name particles every one " - "of them has" - if len(counts) > 1 else "" - ) - raise ValueError( - f"me_frame particle {index} out of range: subprocess " - f"{process} has {len(pdg_list)} external particles{mixed}" - ) - if len(me_frame) == 1: - # A massless particle has no rest frame: its momentum squares to - # zero, so the boost divides by it and every momentum comes back as - # a NaN, silently. Name the mistake instead. - massless = [pdg_list[me_frame[0] - 1] for pdg_list in pdgs - if self.get_mass(pdg_list[me_frame[0] - 1]) == 0.] - if massless: - raise ValueError( - f"me_frame = {me_frame} asks for the rest frame of particle " - f"{me_frame[0]}, which is massless (pdg " - f"{sorted(set(massless))}) and has none. Name a massive " - "particle, or several particles whose sum is massive." - ) - self.me_frame = me_frame - self.incoming_count = 1 if self.is_decay else 2 - - 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 and self.merged_subprocess_data is not None: - # merged_subprocesses.json indexes the *unfiltered* subprocess - # list, so dropping entries here would silently shift every index - # it holds. Refuse rather than mis-map. - raise ValueError( - "merge_subprocesses is not supported for a decay directory " - "with kinematically closed channels" - ) - 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"]) - - def init_event_dir(self) -> None: - output_dir = self.run_card["gridpack"]["output_dir"] - temp_dir = self.run_card["gridpack"]["temp_output_dir"] - if output_dir: - self.run_path = os.path.abspath(os.path.expanduser(output_dir)) - os.makedirs(self.run_path, exist_ok=True) - else: - run_name = self.run_card["run"]["run_name"] - os.makedirs("Events", exist_ok=True) - run_dir_prefix = os.path.join("Events", f"{run_name}_") - existing_run_dirs = glob.glob(f"{run_dir_prefix}*") - run_index = 1 - for run_dir in existing_run_dirs: - run_index_str = run_dir[len(run_dir_prefix):] - if run_index_str.isnumeric(): - run_index = max(run_index, int(run_index_str) + 1) - while True: - try: - self.run_path = f"{run_dir_prefix}{run_index:02d}" - os.mkdir(self.run_path) - break - except FileExistsError: - run_index += 1 - # temporary npy files of the channel generators - if temp_dir: - self.temp_path = os.path.abspath(os.path.expanduser(temp_dir)) - os.makedirs(self.temp_path, exist_ok=True) - else: - self.temp_path = self.run_path - # Absolute on purpose. StatusFile is a C++ object that finishes its - # write from its destructor (rename info.json.tmp -> info.json), and - # that destructor runs whenever Python gets round to collecting the - # process object. When the launcher runs in MG5's process, an - # interrupted run is collected *after* the launch command has restored - # the working directory, and a relative path would then resolve - # somewhere else -- the rename throws a C++ filesystem_error out of a - # destructor, which is an immediate abort. - self.status_file = ms.StatusFile( - os.path.abspath(os.path.join(self.run_path, "info.json"))) - - def init_context(self) -> None: - device_names = self.run_card["run"]["device"] - self.contexts = [] - self.device_types = [] - self.devices = [] - self.pool_sizes = [] - for device_name in device_names: - if ":" in device_name: - device_type, device_index_str = device_name.split(":") - device_index = int(device_index_str) - else: - device_type = device_name - device_index = 0 - self.device_types.append(device_type) - if device_type == "cuda": - device = ms.cuda_device(device_index) - pool_size = self.run_card["run"]["gpu_thread_pool_size"] - elif device_type == "hip": - device = ms.hip_device(device_index) - pool_size = self.run_card["run"]["gpu_thread_pool_size"] - else: - device = ms.cpu_device() - pool_size = self.run_card["run"]["cpu_thread_pool_size"] - self.devices.append(device) - self.pool_sizes.append(pool_size) - self.contexts.append(ms.Context(device=device, thread_count=pool_size)) - - CUSTOM_FUNCTION_KINDS = ("cut", "scale", "pdf", "bias") - - def init_custom_functions(self) -> None: - """Load the functions of [custom_functions], building the libraries - first. They are loaded before the matrix elements, so that they have the - same index in every context and in a gridpack made from this run.""" - self.custom_functions = {} - self.custom_function_paths = [] - for kind in self.CUSTOM_FUNCTION_KINDS: - path = self.run_card["custom_functions"][kind] - if not path: - continue - build_dir = f"custom_{kind}" - if path.endswith(".so") and os.path.isfile(os.path.join(build_dir, "Makefile")): - self.build_custom_function(build_dir) - if not os.path.isfile(path): - raise RuntimeError( - f"[custom_functions] {kind} = {path!r}: file not found. " - "Create a .so library or a .py file with bin/init_custom_func." - ) - apis = load_custom_function(ms, self.contexts, self.device_types, kind, path) - if len({api.index() for api in apis}) != 1: - raise RuntimeError(f"custom {kind} function loaded at different indices") - self.custom_functions[kind] = apis[0] - self.custom_function_paths.append([kind, path]) - logger.info("Using the custom %s function %s", kind, path) - - def init_surrogate(self) -> None: - """Check the [surrogate] settings. The surrogates are loaded per - subprocess, by MadgraphSubprocess.""" - settings = self.run_card["surrogate"] - self.surrogate = None - if not settings["function"]: - return - path = settings["function"] - if path.endswith(".py") and not os.path.isfile(path): - raise RuntimeError(f"[surrogate] function = {path!r}: file not found") - if self.run_card["phasespace"]["sde_strategy"] != "denominators": - raise RuntimeError( - "[surrogate] needs [phasespace] sde_strategy = \"denominators\", " - "as the surrogate does not return channel weights" - ) - if self.run_card["run"]["dummy_matrix_element"]: - raise RuntimeError("[surrogate] cannot be used with dummy_matrix_element") - if self.run_card["gridpack"]["save_gridpack"]: - raise RuntimeError("[surrogate] does not support gridpacks yet") - self.surrogate = dict(settings) - if not self.surrogate["param_card"]: - self.surrogate["param_card"] = self.param_card_path - logger.info( - "Using the surrogate %s, %s with the full matrix element", - path, - "unweighting" if settings["mode"] == "unweight" else "reweighting", - ) - - def weighting_groups(self, subprocess_channels: list[list[int]]) -> list: - """One ms.WeightingGroup per subprocess with a [surrogate]: the second - pass with the full matrix element over its channels, given by their - indices in the event generator.""" - if self.surrogate is None: - return [] - gen_args = self.run_card["generation"] - return [ - ms.WeightingGroup( - channel_indices, - [ms.WeightingStep( - subproc.step_matrix_elements, - mode=self.surrogate["mode"], - max_overweight_truncation=gen_args["max_overweight_truncation"], - freeze_max_weight_after=gen_args["freeze_max_weight_after"], - )], - ) - for subproc, channel_indices in zip(self.subprocesses, subprocess_channels) - if channel_indices - ] - - @staticmethod - def build_custom_function(build_dir: str) -> None: - """Run make in a custom_ directory, logging to compile.log.""" - log_path = os.path.join(build_dir, "compile.log") - with open(log_path, "w") as log: - returncode = subprocess.call( - ["make"], cwd=build_dir, stdout=log, stderr=subprocess.STDOUT - ) - if returncode != 0: - raise RuntimeError( - f"Compilation of {build_dir} failed with exit status {returncode}; " - f"see {log_path} for details" - ) - - def custom_function(self, kind: str, particle_count: int, pid_options: list): - """The ms.CustomFunction of [custom_functions] , or None.""" - api = self.custom_functions.get(kind) - if api is None: - return None - function_type = getattr(ms.CustomFunction, kind) - return ms.CustomFunction(api, function_type, particle_count, pid_options) - - def parse_observable(self, name: str, order_observable: str) -> dict: - parts = name.split("-") - sum_momenta = False - sum_observable = False - ordered = False - multiparticles = self.run_card["multiparticles"] - - if len(parts) == 0: - raise ValueError("Invalid observable name") - elif len(parts) == 1: - # event-level observables - obs_name = parts[0] - select_pids = [] - else: - if parts[-1] == "sum": - sum_observable = True - obs_name = parts[-2] - selection = parts[:-2] - elif parts[-2] == "sum": - sum_momenta = True - obs_name = parts[-1] - selection = parts[:-2] - else: - obs_name = parts[-1] - selection = parts[:-1] - select_pids = [] - order_indices = [] - for mp_name in selection: - mp_parts = mp_name.split("_") - if mp_parts[-1].isnumeric(): - order_indices.append(int(mp_parts[-1])) - select_pids.append(multiparticles["_".join(mp_parts[:-1])]) - ordered = True - else: - order_indices.append(0) - select_pids.append(multiparticles[mp_name]) - - return dict( - observable=obs_name, - select_pids=select_pids, - sum_momenta=sum_momenta, - sum_observable=sum_observable, - order_observable=order_observable if ordered else None, - order_indices=order_indices if ordered else [], - ignore_incoming=True, - name=name, - ) - - def init_cuts(self) -> None: - inf = float("inf") - order_observable = self.run_card["cuts"].get("order_by", "pt") - self.cut_data = [ - CutItem( - observable_kwargs=self.parse_observable(key, order_observable), - min=values.get("min", -inf), - max=values.get("max", inf), - mode=values.get("mode", "all"), - ) - for key, values in self.run_card["cuts"].items() - if key != "order_by" - ] - - # [histograms] key that means "the distribution of the event weight", - # normalised to the unit weight sigma_abs, rather than an observable of the - # momenta: an unweighted sample is a spike at 1 (and at -1 for its negative - # weights) and a partially unweighted one shows its spread. - weight_histogram_key = "weight" - - def init_histograms(self) -> None: - order_observable = self.run_card["histograms"].get("order_by", "pt") - self.hist_data = [] - for key, values in self.run_card["histograms"].items(): - if key == "order_by": - continue - from_weight = key == self.weight_histogram_key - self.hist_data.append(HistItem( - observable_kwargs=( - None if from_weight - else self.parse_observable(key, order_observable)), - min=values["min"], - max=values["max"], - bin_count=values["bin_count"], - name=key, - from_weight=from_weight, - )) - - def ensure_pdf_set(self, pdf_set: str) -> misc.LhapdfPaths: - """Make sure the requested LHAPDF set is available, downloading it into - the first writable PDF directory if needed, and return the (possibly - updated) LHAPDF resolution. Nothing is downloaded when the set is - already present.""" - paths = lhapdf_paths() - if paths.find_set(pdf_set): - return paths - if not paths.config: - logger.debug("no usable lhapdf-config: cannot download %s", pdf_set) - return paths - if not paths.download_path: - logger.debug("no writable PDF directory: cannot download %s", pdf_set) - return paths - try: - from madgraph.interface.common_run_interface import CommonRunCmd - os.makedirs(paths.download_path, exist_ok=True) - logger.info("PDF set %s not found; downloading into %s", - pdf_set, paths.download_path) - CommonRunCmd.install_lhapdf_pdfset_static( - paths.config, paths.download_path, pdf_set) - except Exception as err: - logger.warning("Could not download PDF set %s: %s", pdf_set, err) - return paths - global _LHAPDF - _LHAPDF = paths.with_data_path(paths.download_path) - return _LHAPDF - - def _no_pdf_message(self, pdf_set: str) -> str: - """Explain a missing PDF set in terms of the settings that fix it.""" - searched = os.pathsep.join(self.lhapdf.data_paths) or \ - "(no PDF set directory could be located)" - if self.lhapdf.config: - found = "using lhapdf-config: %s" % self.lhapdf.config - else: - found = ("no usable lhapdf-config was found, so the set could not " - "be downloaded automatically") - return ( - "LHAPDF set %r, requested by [beam] pdf in Cards/run_card.toml, " - "was not found.\n" - " searched: %s\n" - " %s\n" - "Fix this in one of the following ways:\n" - " * install LHAPDF for this MadGraph: MG7> install lhapdf6\n" - " * point MadGraph at an existing one:\n" - " MG7> set lhapdf /path/to/lhapdf-config\n" - " MG7> save options\n" - " * or set it for this directory only, by adding\n" - " lhapdf = /path/to/lhapdf-config\n" - " to %s\n" - " * or set $LHAPDF_DATA_PATH to a directory containing %r." - % (pdf_set, searched, found, - os.path.join("Cards", "me5_configuration.txt"), pdf_set)) - - def init_beam(self) -> None: - beam_args = self.run_card["beam"] - - # built lazily by build_systematics (needs the subprocess data) - self.systematics = None - self.systematics_data = None - self.systematics_context = None - self.event_histograms = None - self.event_histograms_context = None - - 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, - "transverse_mass": ms.EnergyScale.transverse_mass, - "half_transverse_mass": ms.EnergyScale.half_transverse_mass, - "partonic_energy": ms.EnergyScale.partonic_energy, - } - if beam_args["dynamical_scale_choice"] in dynamical_scales: - dynamical_scale_type = dynamical_scales[beam_args["dynamical_scale_choice"]] - else: - raise ValueError("Unknown dynamical scale choice") - self.scale_kwargs = dict( - dynamical_scale_type=dynamical_scale_type, - ren_scale_fixed=beam_args["fixed_ren_scale"], - fact_scale_fixed=beam_args["fixed_fact_scale"], - ren_scale=beam_args["ren_scale"], - fact_scale1=beam_args["fact_scale1"], - fact_scale2=beam_args["fact_scale2"], - scale_factor=float(beam_args.get("scale_factor", 1.)), - ) - if self.is_decay and self.custom_functions.pop("scale", None) is not None: - logger.warning("custom scale ignored: a decay uses the decaying mass") - if self.is_decay and self.custom_functions.pop("pdf", None) is not None: - logger.warning("custom pdf ignored: a decay has no parton densities") - 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, - scale_factor=1., - ) - self.pdf_grid = None - self.lhapdf = None - self.pdf_dir = None - self.fixed_alphas_info = self.write_fixed_alphas_info() - self.alphas_grid = ms.AlphaSGrid(self.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.pdf_set = pdf_set - self.fixed_alphas_info = None - self.lhapdf = self.ensure_pdf_set(pdf_set) - self.pdf_dir = self.lhapdf.find_set(pdf_set) - if self.pdf_dir is None: - raise RuntimeError(self._no_pdf_message(pdf_set)) - self.alphas_grid = ms.AlphaSGrid(os.path.join(self.pdf_dir, pdf_set, f"{pdf_set}.info")) - if "pdf" in self.custom_functions: - # alpha_s still comes from the PDF set of the run card - self.pdf_grid = ms.CustomPdf(self.custom_functions["pdf"]) - else: - self.pdf_grid = ms.PdfGrid( - os.path.join(self.pdf_dir, pdf_set, f"{pdf_set}_0000.dat") - ) - for context in self.contexts: - if not isinstance(self.pdf_grid, ms.CustomPdf): - self.pdf_grid.initialize_globals(context) - self.alphas_grid.initialize_globals(context) - self.running_coupling = ms.RunningCoupling(self.alphas_grid) - - # ------------------------------------------------------------------ - # scale / PDF systematics - # ------------------------------------------------------------------ - def systematics_enabled(self) -> bool: - """The native scale/PDF weights are requested ([systematics] enable, or - the legacy [generation] systematics alias).""" - try: - enabled = bool(self.run_card["systematics"]["enable"]) - except Exception: - enabled = False - if not enabled: - try: - enabled = bool(self.run_card["generation"]["systematics"]) - except Exception: - enabled = False - if enabled and "pdf" in getattr(self, "custom_functions", {}): - # the variations re-evaluate the PDF set, which a custom PDF replaces - if not getattr(self, "_warned_custom_pdf_systematics", False): - logger.warning("systematics are disabled with a custom pdf") - self._warned_custom_pdf_systematics = True - return False - return enabled - - @staticmethod - def pdf_set_info(info_path: str) -> dict: - """Read SetIndex, NumMembers, ErrorType and SetDesc from an LHAPDF - .info file (missing keys get neutral defaults).""" - info = {"SetIndex": 0, "NumMembers": 1, "ErrorType": "", "SetDesc": ""} - try: - with open(info_path) as f: - for line in f: - key, sep, value = line.partition(":") - key = key.strip() - if not sep or key not in info: - continue - value = value.strip().strip('"').strip("'") - if key in ("SetIndex", "NumMembers"): - info[key] = int(value) - else: - info[key] = value - except OSError as err: - logger.warning("could not read %s: %s", info_path, err) - return info - - def resolve_pdf_set_name(self, entry: str) -> tuple[str, int | None]: - """Turn a [systematics] pdf entry ("", "", "@", - "@") into (set name, member or None) using the LHAPDF - pdfsets.index when an id is given.""" - member = None - if "@" in entry: - entry, member_str = entry.split("@", 1) - member = int(member_str) - entry = entry.strip() - if not entry.isdigit(): - return entry, member - lhaid = int(entry) - index_path = os.path.join(self.pdf_dir or "", "pdfsets.index") - sets = [] - try: - with open(index_path) as f: - for line in f: - parts = line.split() - if len(parts) >= 2 and parts[0].isdigit(): - sets.append((int(parts[0]), parts[1])) - except OSError: - raise ValueError( - "cannot resolve LHAPDF id %s: no %s" % (lhaid, index_path)) - for set_id, name in sets: - if set_id == lhaid: - return name, member - # an id inside a set's member range - sets.sort() - for (set_id, name), (next_id, _) in zip(sets, sets[1:] + [(None, None)]): - if set_id < lhaid and (next_id is None or lhaid < next_id): - return name, lhaid - set_id if member is None else member - raise ValueError("cannot resolve LHAPDF id %s" % lhaid) - - def resolve_pdf_variations(self) -> list: - """Expand the [systematics] pdf entries into ms.PdfMemberSpec objects - (downloading the sets when needed). The nominal member itself is never - listed: the calculator adds it as the central weight of its group.""" - entries = self.run_card["systematics"]["pdf"] - if isinstance(entries, str): - entries = [entries] - specs = [] - seen = set() - for raw in entries: - entry = str(raw).strip() - if not entry or entry.lower() in ("none", "no", "false"): - continue - if entry.lower() == "central": - continue - if entry.lower() == "errorset": - set_name, member = self.pdf_set, None - else: - set_name, member = self.resolve_pdf_set_name(entry) - variation_lhapdf = self.ensure_pdf_set(set_name) - variation_dir = variation_lhapdf.find_set(set_name) - if variation_dir is None: - logger.warning("PDF set %s not available; its variation is skipped", set_name) - continue - info_path = os.path.join(variation_dir, set_name, f"{set_name}.info") - if not os.path.exists(info_path): - logger.warning("PDF set %s not available; its variation is skipped", set_name) - continue - info = self.pdf_set_info(info_path) - is_nominal = set_name == self.pdf_set - if member is not None: - members = [member] - else: - members = list(range(info["NumMembers"])) - if is_nominal: - members = members[1:] - if not members: - logger.warning( - "PDF set %s has a single member: no PDF uncertainty can be " - "computed from it. Use a set with error members in " - "[systematics] pdf (e.g. a PDF4LHC or NNPDF error set).", - set_name) - for m in members: - if is_nominal and m == 0: - continue - if (set_name, m) in seen: - continue - seen.add((set_name, m)) - grid_file = os.path.join(variation_dir, set_name, f"{set_name}_{m:04d}.dat") - if not os.path.exists(grid_file): - logger.warning("PDF member file %s not found; skipped", grid_file) - continue - specs.append(ms.PdfMemberSpec( - set_name=set_name, set_lhaid=info["SetIndex"], member=m, - grid_file=grid_file, info_file=info_path, - error_type=info["ErrorType"], description=info["SetDesc"], - )) - return specs - - def build_systematics_args(self) -> list: - """One ms.SubprocessSystArgs per (unmerged) subprocess: the alpha_s power - of |M|^2 and the beam parton ids of every flavor index, indexed like the - LHECompleter (subprocess_index / flavor_index of the combined events).""" - args = [] - for meta in self.subprocess_data: - beam_pdgs = [ - [int(flavor["options"][0][0]), int(flavor["options"][0][1])] - for flavor in meta["flavors"] - ] - args.append(ms.SubprocessSystArgs( - qcd_power=int(meta.get("qcd_power", -1)), beam_pdgs=beam_pdgs)) - return args - - def build_systematics(self): - """Create the ms.SystematicsCalculator from the [systematics] section, or - None when disabled. Also keeps the JSON description needed by gridpacks.""" - if not self.systematics_enabled(): - self.systematics = None - return None - syst = self.run_card["systematics"] - config = ms.SystematicsConfig() - config.mur = [float(v) for v in syst["mur"]] - config.muf = [float(v) for v in syst["muf"]] - config.together = bool(syst["together"]) - config.dyn_scales = self.resolve_dynamical_scales() - # alternative dynamical scales get the same factor as the nominal one - config.scale_factor = float(self.run_card["beam"].get("scale_factor", 1.)) - config.write_inputs = bool(syst["write_inputs"]) - config.has_pdf = not self.leptonic - # Without parton luminosity -- a decay, or leptonic beams -- there is no - # nominal PDF *set* to describe and no PDF variation to compute, so - # has_pdf gates the nominal-set metadata and the member list below. - # The alpha_s grid is a separate matter: the mu_R variations always need - # one, so its .info file has to be recorded whatever has_pdf says. - # init_beam() skips the PDF lookup for a decay only -- leptonic beams - # still resolve a set and take alpha_s from it -- so the path follows - # self.pdf_dir, and falls back to the constant grid written for a decay. - if self.pdf_dir is not None: - info_path = os.path.join(self.pdf_dir, self.pdf_set, f"{self.pdf_set}.info") - else: - info_path = self.fixed_alphas_info - if config.has_pdf: - info = self.pdf_set_info(info_path) - config.nominal_set_name = self.pdf_set - config.nominal_lhaid = info["SetIndex"] - config.nominal_error_type = info["ErrorType"] - config.nominal_description = info["SetDesc"] - config.pdf_members = self.resolve_pdf_variations() - else: - config.pdf_members = [] - args = self.build_systematics_args() - # the PDFs, alpha_s and (for mixed-order subprocesses) the matrix - # elements are evaluated with the batched madspace functions on this - # CPU context - self.systematics_context = ms.Context(device=ms.cpu_device(), thread_count=1) - matrix_elements, flavor_remap, me_backend = self.build_systematics_matrix_elements( - args, self.systematics_context) - self.systematics = ms.SystematicsCalculator( - config, args, self.pdf_grid, self.alphas_grid, - self.systematics_context, matrix_elements, flavor_remap) - for warning in self.systematics.warnings: - logger.warning("systematics: %s", warning) - self.systematics_data = { - "config": json.loads(config.to_json()), - "subproc_args": [json.loads(a.to_json()) for a in args], - "nominal_grid_file": os.path.join( - self.pdf_dir, self.pdf_set, f"{self.pdf_set}_0000.dat") - if config.has_pdf else None, - "nominal_info_file": info_path, - # matrix elements re-evaluated for the mixed-order subprocesses - "me_paths": [meta["me_path"] for meta in self.subprocess_data], - "me_backend": me_backend, - "flavor_remap": flavor_remap, - "me_frame": self.me_frame, - "incoming_count": self.incoming_count, - } - return self.systematics - - _DYN_SCALE_CODES = { - "transverse_energy": 1, "transverse_mass": 2, - "half_transverse_mass": 3, "partonic_energy": 4, - } - - def resolve_dynamical_scales(self) -> list: - """[systematics] dynamical_scale entries (names or systematics.py codes - 1-4) as codes, without the choice the events were generated with.""" - try: - entries = self.run_card["systematics"]["dynamical_scale"] - except Exception: - return [] - if isinstance(entries, str): - entries = [entries] - beam = self.run_card["beam"] - generated = None - if not (beam["fixed_ren_scale"] and beam["fixed_fact_scale"]): - generated = self._DYN_SCALE_CODES.get(beam["dynamical_scale_choice"]) - codes = [] - for raw in entries: - entry = str(raw).strip().lower() - if not entry: - continue - if entry.isdigit(): - code = int(entry) - elif entry in self._DYN_SCALE_CODES: - code = self._DYN_SCALE_CODES[entry] - else: - raise ValueError("unknown dynamical scale choice %r in [systematics] " - "dynamical_scale" % raw) - if code not in (1, 2, 3, 4): - raise ValueError("dynamical scale code %s must be 1-4" % code) - if code == generated or code in codes: - continue - codes.append(code) - return codes - - def build_systematics_matrix_elements(self, args, context): - """For the subprocesses whose |M|^2 mixes several powers of alpha_s - (qcd_power = -1), load the matrix element into the systematics CPU - context so that the renormalisation scale variations can re-evaluate it - at the varied alpha_s. Returns ([MatrixElement or None per subprocess], - flavor remap, backend name); empty when nothing is needed or no CPU - library exists.""" - need = [i for i, a in enumerate(args) if a.qcd_power < 0] - empty = ([], [], None) - if not need or self.run_card["run"]["dummy_matrix_element"]: - return empty - cpu_backends = [b for b in self.backends if not b.startswith(("cuda", "hip"))] - if not cpu_backends: - logger.warning("systematics: no CPU matrix-element library available, the " - "renormalisation scale variations of the mixed-order " - "subprocesses cannot be computed") - return empty - backend = cpu_backends[0] - matrix_elements = [] - flavor_remap = [] - for i, meta in enumerate(self.subprocess_data): - flavor_remap.append([int(f["index"]) for f in meta["flavors"]]) - if i not in need: - matrix_elements.append(None) - continue - api = context.load_matrix_element( - meta["me_path"].format(device=backend), self.param_card_path, - me_parameters(self.run_card)) - matrix_elements.append(ms.MatrixElement( - api, - [ms.MatrixElement.momenta_in, ms.MatrixElement.alpha_s_in, - ms.MatrixElement.flavor_in], - [ms.MatrixElement.matrix_element_out], - False, - self.me_frame, - self.incoming_count, - )) - logger.info("systematics: re-evaluating the matrix element of %d mixed-order " - "subprocess(es) for the renormalisation scale variations", len(need)) - return matrix_elements, flavor_remap, backend - - def build_event_histograms(self): - """ms.EventHistograms for the [histograms] observables, filled with the - written events and all their variation weights (info.json - "event_histograms"); None without histograms.""" - if not self.hist_data \ - or not self.run_card["run"]["postprocessing_histograms"]: - self.event_histograms = None - return None - context = ms.Context(device=ms.cpu_device(), thread_count=1) - specs = [ - ms.EventHistogramSpec( - name=item.name, min=item.min, max=item.max, - bin_count=item.bin_count, from_weight=item.from_weight) - for item in self.hist_data - ] - from_momenta = [item for item in self.hist_data if not item.from_weight] - observables = [] - for meta in self.subprocess_data: - if not from_momenta: - # nothing to evaluate on the momenta (only the weight is - # histogrammed): no observable runtime to build - observables.append(None) - continue - all_pids = clean_pids(meta["incoming"]) + clean_pids(meta["outgoing"]) - values = ms.ObservableValues([ - build_observable( - observable_pids(meta, item.observable_kwargs), - len(meta["incoming"]), item.observable_kwargs - ) - for item in from_momenta - ]) - observables.append(ms.SubprocessObservables(values, len(all_pids))) - # the weight histograms are drawn in units of the unit weight sigma_abs - # (the cross section of a sample without negative weights) - self.event_histograms_context = context - self.event_histograms = ms.EventHistograms( - context, specs, observables, - reference_weight=self._mean_event_weight()) - return self.event_histograms - - def _mean_event_weight(self) -> float: - """The unit weight of the unweighted events, sigma_abs, used as the - unit of the weight histograms: the cross section of a positive sample, - and for one with negative weights (an interference) what puts them at - +-1 rather than dividing by a signed, possibly small, total. 0 (the raw - weight) when it is not available.""" - try: - return float(self.event_generator.status().mean_abs) - except Exception as error: - logger.debug("no cross section for the weight histograms: %s", error) - return 0. - - def log_systematics_summary(self) -> None: - """Print the scale/PDF uncertainties on the total cross section (the - per-variation cross sections are in events.weights.json / info.json).""" - if self.systematics is None or self.systematics.weight_count == 0: - return - summary = json.loads(self.systematics.summary()) - # the percentages come from systematics_summary, which get_result() - # reads too: the box and the scan row must not disagree, so neither - # computes them itself - xsec = systematics_summary.nominal_cross_section(summary) - if xsec is None: - return - def format_variation(up, down): - # right-justify the signed numbers (not the sign alone) so the % - # values line up across rows, with extra spacing between the two - return f"{'+%.3g' % up:>6}% {'-%.3g' % down:>6}%" - - scale_count = sum(1 for v in self.systematics.variations if v.is_scale) - member_count = len(self.systematics.members) - rows = [("Variations per event:", - f"{self.systematics.weight_count} ({scale_count} scale, {member_count} PDF members)")] - for pdf in summary.get("pdf", []): - rows.append(("PDF set:", f"{pdf['pdf_set']}, {pdf['error_type']}")) - rows.append(("Original cross-section:", f"{xsec} pb")) - scale = systematics_summary.scale_percentages(summary) - if scale is not None: - rows.append(("Scale variation:", format_variation(*scale))) - for _entry, up, down in systematics_summary.pdf_percentages(summary): - rows.append(("PDF variation:", format_variation(up, down))) - if self.event_generator_config.verbosity == ms.Verbosity.pretty: - box = ms.PrettyBox("Systematics", len(rows), [24, 0]) - box.set_column(0, [label for label, _ in rows]) - box.set_column(1, [value for _, value in rows]) - box.print_first() - else: - for label, value in rows: - logger.info("systematics, %s %s", label, value) - - 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 needs_systematics_matrix_elements(self) -> bool: - """The systematics will re-evaluate a matrix element per scale point. - - True when the native weights are on and some subprocess mixes alpha_s - powers (qcd_power < 0), i.e. when build_systematics_matrix_elements() - loads a matrix element. Conservative: it does not check that a CPU - backend exists, so it can be True where that call ends up loading - nothing. - """ - if not self.systematics_enabled(): - return False - if self.run_card["run"]["dummy_matrix_element"]: - return False - return any(int(meta.get("qcd_power", -1)) < 0 - for meta in self.subprocess_data) - - def init_generator_config(self) -> None: - run_args = self.run_card["run"] - gen_args = self.run_card["generation"] - vegas_args = self.run_card["vegas"] - cfg = ms.GeneratorConfig() - cfg.target_count = gen_args["events"] - cfg.vegas_damping = vegas_args["damping"] - cfg.max_overweight_truncation = gen_args["max_overweight_truncation"] - cfg.freeze_max_weight_after = gen_args["freeze_max_weight_after"] - if self.surrogate is not None and self.surrogate["mode"] == "unweight": - # [generation] applies to the second pass, which gives the events - cfg.max_overweight_truncation = ( - self.surrogate["first_pass_max_overweight_truncation"] - ) - cfg.freeze_max_weight_after = ( - self.surrogate["first_pass_freeze_max_weight_after"] - ) - cfg.start_batch_size = vegas_args["start_batch_size"] - cfg.max_batch_size = vegas_args["max_batch_size"] - cfg.survey_min_iters = gen_args["survey_min_iters"] - cfg.survey_max_iters = gen_args["survey_max_iters"] - cfg.survey_target_precision = gen_args["survey_target_precision"] - cfg.optimization_patience = vegas_args["optimization_patience"] - cfg.optimization_threshold = vegas_args["optimization_threshold"] - cfg.cpu_batch_size = gen_args["cpu_batch_size"] - cfg.gpu_batch_size = gen_args["gpu_batch_size"] - cfg.verbosity = resolve_verbosity(run_args["verbosity"]) - cfg.combine_thread_count = run_args["combine_thread_pool_size"] - if self.needs_systematics_matrix_elements() and cfg.combine_thread_count != 1: - # SystematicsCalculator::compute() runs inside the combine workers, - # and a matrix element's per-thread process instances live in a - # ThreadResource indexed by ThreadPool::thread_index() -- the - # *calling* thread's index. A parallel combine pool would therefore - # need one process instance per worker; building N of them to index - # them safely is not worth it, the more so as - # SystematicsCalculator::matrix_elements() already serialises every - # evaluation on its own mutex, so that parallelism was not being - # used anyway. Run the combine stage single-threaded instead. - logger.info( - "systematics: the combine stage runs single-threaded because a " - "subprocess needs its matrix element re-evaluated for the " - "renormalisation scale variations") - cfg.combine_thread_count = 1 - cfg.cut_efficiency_threshold = gen_args["cut_efficiency_threshold"] - cfg.max_cut_repetitions = gen_args["max_cut_repetitions"] - self.event_generator_config = cfg - self.event_generator = None - - def init_subprocesses(self) -> None: - self.backends = self.compile_matrix_elements() - self.subprocesses = [] - if self.merged_subprocess_data is None: - for subproc_id, meta in enumerate(self.subprocess_data): - self.subprocesses.append(MadgraphSubprocess(self, meta, subproc_id)) - else: - for subproc_id, meta in enumerate(self.merged_subprocess_data): - self.subprocesses.append( - MadgraphSubprocess(self, meta, subproc_id, self.subprocess_data) - ) - - def compile_matrix_elements(self) -> list[str]: - """Build the matrix-element library of every subprocess, and return the - build backend of each requested device (in the same order). - - A 'cpu' device builds the backend named by the 'cpu_mode' entry; the - cuda/hip devices ignore cpu_mode and build their own backend. A - cpu_mode of 'auto' is resolved to the actual SIMD backend chosen on - this machine. - - SubProcesses/makefile is a dispatcher over the P* directories, so a - single 'make -j N' there builds all the subprocesses at once with one - shared pool of N jobs: no subprocess is built with N jobs while the - others wait, and none is limited to N/#subprocesses jobs either. - """ - device_names = self.run_card["run"]["device"] - if not isinstance(device_names, list): - device_names = [device_names] - cpu_mode = self.run_card["run"]["cpu_mode"] - backends = [backend_of(name, cpu_mode) for name in device_names] - if not self.subprocess_data: - return backends - - first_proc_path = self.subprocess_data[0]["path"] - subproc_path = os.path.dirname(first_proc_path) - - fptype = fptype_of(self.run_card["run"]["precision"]) - - # Resolve cpu_mode='auto' once (the build rules pick the best SIMD - # backend available here), so all subprocesses agree on the library names. - auto_backend = None - if "auto" in backends: - auto_backend = resolve_auto_backend(first_proc_path) - logger.info("cpu_mode 'auto' resolved as '%s'", auto_backend) - resolved = [ - auto_backend if backend == "auto" else backend - for backend in backends - ] - - nb_core = self.run_card["run"]["cpu_thread_pool_size"] - if not nb_core or nb_core < 0: - nb_core = os.cpu_count() or 1 - - log_path = os.path.join(self.run_path, "compile_subprocesses.log") - for backend in resolved: - missing = [ - meta for meta in self.subprocess_data - if not os.path.isfile(meta["me_path"].format(device=backend)) - ] - if not missing: - continue - logger.info( - f"Start compilation of SubProcesses for device '{backend}' " - f"({len(missing)} subprocess(es), {nb_core} parallel job(s)), " - f"see log detail in {log_path}" - ) - start_time = time.time() - self.make_subprocesses( - subproc_path, - [f"BACKEND={backend}", f"FPTYPE={fptype}", "USEBUILDDIR=1"], - nb_core, log_path, - ) - logger.info( - f"Compilation of SubProcesses done in {time.time() - start_time:.1f} s" - ) - return resolved - - @staticmethod - def make_subprocesses( - subproc_path: str, args: list[str], nb_core: int, log_path: str - ) -> None: - """Run 'make -j nb_core' in SubProcesses/, appending the full build - output to log_path (one run per device, so the file is appended to).""" - command = ["make", f"-j{nb_core}"] + args - with open(log_path, "a") as log: - log.write(f"\n$ cd {subproc_path} && {' '.join(command)}\n") - log.flush() - returncode = subprocess.call( - command, cwd=subproc_path, stdout=log, stderr=subprocess.STDOUT - ) - if returncode != 0: - raise RuntimeError( - f"Compilation of the SubProcesses failed with exit status " - f"{returncode}; see {log_path} for details" - ) - - def build_event_generator(self, phasespaces: list[PhaseSpace]) -> ms.EventGenerator: - channel_generators = [] - subprocess_channels = [[] for _ in self.subprocesses] - for i, (subproc, phasespace) in enumerate(zip(self.subprocesses, phasespaces)): - for integrand, channel in zip( - subproc.build_integrands(phasespace), phasespace.channels - ): - if channel.event_generator is None: - channel.event_generator = ms.ChannelEventGenerator( - contexts=self.contexts, - integrand=integrand, - event_file=os.path.join(self.temp_path, f"events.{i}.{channel.name}.npy"), - weight_file=os.path.join(self.temp_path, f"weights.{i}.{channel.name}.npy"), - config=self.event_generator_config, - subprocess_index=i, - name=f"{i}.{channel.name}", - histograms=subproc.histograms - ) - subprocess_channels[i].append(len(channel_generators)) - channel_generators.append(channel.event_generator) - - event_generator = ms.EventGenerator( - contexts=self.contexts, - channels=channel_generators, - status_file=self.status_file, - config=self.event_generator_config, - seed=self.run_seed, - groups=self.weighting_groups(subprocess_channels), - ) - unused_globals = ( - set(self.contexts[0].global_names()) - event_generator.used_globals() - ) - for context in self.contexts: - for global_name in unused_globals: - context.delete_global(global_name) - return event_generator - - def survey_phasespaces( - self, phasespaces: list[PhaseSpace | None], survey_pass: int = 0 - ) -> ms.EventGenerator | None: - ps_filtered = [ps for ps in phasespaces if ps is not None] - if len(ps_filtered) == 0: - return None - event_generator = self.build_event_generator(ps_filtered) - event_generator.survey(survey_pass) - return event_generator - - def survey(self) -> None: - # survey_pass distinguishes the survey() calls below: "both" mode can - # re-survey a channel carried over unchanged from the multichannel pass - # into the final (simplified) pass, and both passes schedule jobs on the - # same underlying ChannelEventGenerator. The explicit pass index keeps each - # pass's job seeds independent of the other passes' job counts, rather than - # depending on call history. - 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 -- and 'both'/'auto' end up - # simplifying towards it. 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 in ["multichannel", "both", "auto"]: - self.phasespaces = [ - subproc.build_multichannel_phasespace() - for subproc in self.subprocesses - ] - self.event_generator = self.survey_phasespaces(self.phasespaces, 0) - elif phasespace_mode == "flat": - self.phasespaces = [ - subproc.build_flat_phasespace() - for subproc in self.subprocesses - ] - self.event_generator = self.survey_phasespaces(self.phasespaces, 0) - else: - raise ValueError("Unknown phasespace mode") - - channel_status = self.event_generator.channel_status() - chan_offset = 0 - madnis_enabled = False - for subproc, ps in zip(self.subprocesses, self.phasespaces): - mean = 0. - variance = 0. - count_opt = 0 - for status in channel_status[chan_offset:chan_offset + len(ps.channels)]: - mean += status.mean_abs - variance += status.error_abs**2 - count_opt += status.count_opt - # A subprocess whose matrix element is identically zero (an FCNC - # channel with every Wilson coefficient at zero, say) integrates to - # exactly zero, so there is no relative spread to speak of. Report 0 - # rather than dividing by it: that asks for the smallest networks and - # leaves MadNIS off, which is what an empty subprocess wants. - rsd = (variance * count_opt)**0.5 / mean if mean else 0. - subproc.set_madnis_auto_settings(rsd) - chan_offset += len(ps.channels) - if subproc.madnis_settings["enable"]: - madnis_enabled = True - if not ( - phasespace_mode == "both" or (phasespace_mode == "auto" and madnis_enabled) - ): - return - - phasespaces_multi = self.phasespaces - cross_sections = [] - index = 0 - for phasespace in phasespaces_multi: - channel_count = len(phasespace.channels) - # rank the channels by the integral of |w|: a channel whose signed - # integral cancels (an interference) can still carry most of the - # variance, and folding it into the flat channel would blow it up - cross_sections.append([ - status.mean_abs - for status in channel_status[index:index + channel_count] - ]) - index += channel_count - - self.phasespaces = [ - subproc.simplify_phasespace(ps_multi, cross_secs) - for subproc, ps_multi, cross_secs in zip( - self.subprocesses, phasespaces_multi, cross_sections - ) - ] - if any( - ps_multi is not ps_both - for ps_multi, ps_both in zip(phasespaces_multi, self.phasespaces) - ): - # distinct survey_pass: a channel carried over unchanged from the - # multichannel pass (pass 0) into this resurvey must not share its - # seed stream with that earlier pass. - self.event_generator = self.survey_phasespaces(self.phasespaces, 1) - - def train_madnis(self) -> None: - madnis_args = self.run_card["madnis"] - if not any(subproc.madnis_settings["enable"] for subproc in self.subprocesses): - return - - gen_args = self.run_card["generation"] - run_args = self.run_card["run"] - - verbosity = resolve_verbosity(run_args["verbosity"]) - madnis_phasespaces = [] - training_args = [] - self.event_generator = None - for subproc, phasespace in zip(self.subprocesses, self.phasespaces): - for channel in phasespace.channels: - channel.event_generator = None - - config = ms.MadnisConfig() - config.learning_rate = subproc.madnis_settings["lr"] - config.batches = subproc.madnis_settings["train_batches"] - config.log_interval = madnis_args["log_interval"] - config.integration_history_length = madnis_args["integration_history_length"] - config.channel_dropping_interval = madnis_args["channel_dropping_interval"] - config.channel_dropping_threshold = madnis_args["channel_dropping_threshold"] - config.cpu_generator_batch_size = gen_args["cpu_batch_size"] - config.gpu_generator_batch_size = gen_args["gpu_batch_size"] - config.gpu_generator_batch_granularity = madnis_args["gpu_generator_batch_granularity"] - config.generator_target_size_factor = madnis_args["generator_target_size_factor"] - config.batch_size_offset = madnis_args["batch_size_offset"] - config.batch_size_per_channel = subproc.madnis_settings["batch_size_per_channel"] - config.uniform_channel_ratio = madnis_args["uniform_channel_ratio"] - config.lr_schedule = madnis_args["lr_scheduler"] - config.adam_beta1 = madnis_args["adam_beta1"] - config.adam_beta2 = madnis_args["adam_beta2"] - config.adam_eps = madnis_args["adam_eps"] - config.adam_weight_decay = madnis_args["adam_weight_decay"] - config.grad_clip_threshold = madnis_args["grad_clip_threshold"] - config.buffer_capacity = madnis_args["buffer_capacity"] - config.minimum_buffer_size = madnis_args["minimum_buffer_size"] - config.buffered_steps_fraction = madnis_args["buffered_steps_fraction"] - config.buffer_skip_batches = madnis_args["buffer_skip_batches"] - config.buffer_unweighting_quantile = madnis_args["buffer_unweighting_quantile"] - config.fixed_cwnet_fraction = subproc.madnis_settings["fixed_cwnet_fraction"] - config.softclip_threshold = madnis_args["softclip_threshold"] - config.compressed_channel_weight_count = madnis_args["compressed_channel_weight_count"] - phasespace = subproc.build_madnis(phasespace) - madnis_phasespaces.append(phasespace) - training_args.append( - ms.TrainingArgs( - config=config, - integrands=subproc.build_integrands( - phasespace, - madnis_training=True, - drop_cuts_and_rescale=madnis_args["drop_zero_integrands"] - ), - cwnet=phasespace.cwnet, - ) - ) - - gen_context = self.contexts[0] - opt_context = ms.Context( - device=self.devices[0], thread_count=self.pool_sizes[0] - ) - opt_context.copy_globals_from(gen_context) - - madnis_training = ms.MultiMadnisTraining( - generator_context=gen_context, - optimizer_context=opt_context, - training_args=training_args, - verbosity=verbosity, - status_file=self.status_file, - # Reuses the run's resolved seed (also used by build_event_generator()). - # Only the single-channel CPU sample-generation path is currently seeded - # -- buffered training and GPU multi-channel batches are still - # non-deterministic. - seed=self.run_seed, - ) - madnis_training.train() - for phasespace, active_channels in zip( - madnis_phasespaces, madnis_training.active_channels() - ): - phasespace.channels = [ - phasespace.channels[index] for index in active_channels - ] - del madnis_training - del opt_context - self.phasespaces = madnis_phasespaces - for context in self.contexts[1:]: - context.copy_globals_from(self.contexts[0]) - self.event_generator = self.build_event_generator(madnis_phasespaces) - self.madnis_trained = True - - def update_madnis_status_single( - self, batch: int, batch_target: int, loss: float, lr: float, channel_count: int - ) -> None: - now = time.time() - if batch + 1 < batch_target: - if now - self.last_update_time < 0.1: - return - self.last_update_time = now - progress_bar = ms.format_progress((batch + 1) / batch_target, 52) - time_diff = now - self.madnis_wall_time - time_str = f"{format_time(time_diff)}" - else: - progress_bar = "" - wall_diff = now - self.madnis_wall_time - cpu_diff = time.process_time() - self.madnis_cpu_time - time_str = ( - f"{format_time(wall_diff, centi=True)} wall, " - f"{format_time(cpu_diff, centi=True)} cpu" - ) - batch_str = f"{batch + 1} / {batch_target}" - self.madnis_box.set_column(1, [ - f"{batch_str:<15} {progress_bar}", - f"{loss:>.4f}", - f"{channel_count}", - time_str - ]) - self.madnis_box.print_update() - - def update_madnis_status_multi( - self, - subproc_id: int, - batch: int, - batch_target: int, - loss: float, - lr: float, - channel_count: int - ) -> None: - now = time.time() - subproc_count = len(self.subprocesses) - if batch + 1 < batch_target: - if now - self.last_update_time < 0.1: - return - self.last_update_time = now - progress_bar = ms.format_progress((batch + 1) / batch_target, 34) - progress_bar_all = ms.format_progress( - (subproc_id * batch_target + batch + 1) / (subproc_count * batch_target), - 52 - ) - time_str = f"{format_time(now - self.madnis_wall_time)}" - subproc_str = f"{subproc_id} / {subproc_count}" - elif subproc_id < subproc_count - 1: - progress_bar = "" - progress_bar_all = ms.format_progress( - ((subproc_id + 1) * batch_target + 1) / (subproc_count * batch_target), - 52 - ) - time_str = f"{format_time(now - self.madnis_wall_time)}" - subproc_str = f"{subproc_id} / {subproc_count}" - else: - progress_bar = "" - progress_bar_all = "" - wall_diff = now - self.madnis_wall_time - cpu_diff = time.process_time() - self.madnis_cpu_time - time_str = ( - f"{format_time(wall_diff, centi=True)} wall, " - f"{format_time(cpu_diff, centi=True)} cpu" - ) - subproc_str = f"{subproc_count} / {subproc_count}" - batch_str = f"{batch + 1} / {batch_target}" - self.madnis_upper_box.set_column(1, [ - f"{subproc_str:<15} {progress_bar_all}", - time_str, - ]) - self.madnis_lower_box.set_row(subproc_id + 1, [ - f"{subproc_id}", - f"{loss:>.4f}", - f"{channel_count}", - f"{batch_str:<15} {progress_bar}", - ]) - self.madnis_upper_box.print_update() - self.madnis_lower_box.print_update() - - @property - def gridpack_run_mode(self) -> str: - """How this run ends: a gridpack run may skip the events (see the - [gridpack] run_mode run card option); any other run is a regular one.""" - if not self.run_card["gridpack"]["save_gridpack"]: - return "regular" - return self.run_card["gridpack"]["run_mode"] - - def finish_without_events(self) -> None: - """The end of a gridpack run that writes no events: only what the - gridpack needs is prepared, then it is stored.""" - if self.gridpack_run_mode == "minimal": - # a trained MadNIS needs nothing more, a VEGAS grid has to converge - if not self.madnis_trained: - self.event_generator.optimize() - else: - self.event_generator.fix_max_weights() - self.build_systematics() - self.save_gridpack() - - def generate_events(self) -> None: - if self.gridpack_run_mode != "regular": - self.finish_without_events() - return - start_time = get_start_time() - self.event_generator.generate() - output_format = self.run_card["run"]["output_format"] - systematics = self.build_systematics() - histograms = self.build_event_histograms() - if output_format == "compact_npy": - self.lhe_completer = None - self.event_generator.combine_to_compact_npy( - os.path.join(self.run_path, "events.npy"), systematics, histograms - ) - self.write_lhe_header(systematics) - self.save_lhe_completer() - elif output_format == "lhe_npy": - self.lhe_completer = self.build_lhe_completer() - self.event_generator.combine_to_lhe_npy( - os.path.join(self.run_path, "events.npy"), self.lhe_completer, - systematics, histograms - ) - self.write_lhe_header(systematics) - elif output_format == "lhe": - self.lhe_completer = self.build_lhe_completer() - lhe_path = os.path.join(self.run_path, "events.lhe") - self.event_generator.combine_to_lhe( - lhe_path, self.lhe_completer, - self.build_lhe_meta(), systematics, histograms - ) - # Ship the LHE compressed by default. These files are large and - # very compressible, madevent has always stored its events - # gzipped, and every consumer here already accepts either form - # (see _find_event_file). misc.gzip replaces events.lhe with - # events.lhe.gz, and switches to an external multithreaded tool - # above 256 MB. - misc.gzip(lhe_path) - else: - raise ValueError("Unknown output format") - self.write_hwu(histograms) - self.make_plots(histograms) - if systematics is not None: - self.write_systematics_sidecar() - self.log_systematics_summary() - self.save_gridpack() - - def save_lhe_completer(self) -> None: - """Keep the LHE completer next to compact_npy events, so npy_to_lhe - can complete them into an LHE file later (a shower or MadSpin requested - after the run). Never fatal to the run.""" - try: - self.build_lhe_completer().save( - os.path.join(self.run_path, npy_to_lhe.COMPLETER_FILE)) - except Exception as err: - logger.warning("could not save the LHE completer; these events " - "cannot be converted to LHE later: %s", err) - - hwu_file_name = "MADatLO.HwU" - plot_dir_name = "plots" - - def make_plots(self, histograms) -> None: - """Draw the event-sample histograms into Events//plots, with the - scale and PDF bands info.json carries, and say where they landed. - - This does not go through the HwU file: the bands are already computed - per bin by madspace, so drawing them directly is both fewer steps and - more than histograms.py can do without the LHAPDF python module. - """ - if histograms is None or not self.run_card["run"]["make_plots"]: - return - out_dir = os.path.join(self.run_path, self.plot_dir_name) - try: - data = json.loads(histograms.to_json(self.systematics)) - written = plots.render(data, out_dir) - except plots.BackendMissing as error: - logger.info("matplotlib is not available (%s): no plots. The " - "distributions are in %s", - error, os.path.join(self.run_path, "info.json")) - return - except Exception as error: - logger.warning("could not draw the histograms: %s", error) - return - if not written: - return - logger.info("%d plot(s) stored in %s", len(written), out_dir) - - def write_hwu(self, histograms) -> None: - """Write the event-sample histograms next to the events in the HwU - format (MADatLO.HwU), the one an aMC@NLO run writes as MADatNLO.HwU - and madgraph/various/histograms.py reads: an mg7 distribution can then - be overlaid on an NLO one without converting anything in between. The - numbers are the ones info.json carries, bands included, so this is off - by default -- it is a second copy in another format, for the runs that - are going to be plotted.""" - if histograms is None or not self.run_card["run"]["write_hwu"]: - return - try: - data = json.loads(histograms.to_json(self.systematics)) - summary = (json.loads(self.systematics.summary()) - if self.systematics is not None else None) - text = hwu_output.to_hwu(data, summary) - except Exception as error: - logger.warning("could not write the HwU histograms: %s", error) - return - if not text: - return - with open(os.path.join(self.run_path, self.hwu_file_name), "w") as f: - f.write(text) - - def write_systematics_sidecar(self) -> None: - """Describe the variation weights next to the event file - (events.weights.json): ids, scale factors, PDF set/member and the - text, so npy consumers get the same metadata as the LHE - header, plus the per-variation cross sections.""" - if self.systematics is None: - return - data = json.loads(self.systematics.summary()) - data["initrwgt"] = self.systematics.initrwgt() - data["columns"] = ["rwgt_%d" % i for i in self.systematics.weight_ids] - with open(os.path.join(self.run_path, "events.weights.json"), "w") as f: - json.dump(data, f, indent=1) - - @staticmethod - def _histogram_mean(hist_min, hist_max, bin_values): - """Cross-section-weighted mean of a histogrammed observable over its - range. ``bin_values`` carries the underflow and overflow bins first and - last, as both madspace histogram kinds do; they have no position and - are left out.""" - values = list(bin_values)[1:-1] - total = sum(values) - if not values or total == 0: - return None - width = (hist_max - hist_min) / len(values) - return sum(v * (hist_min + (i + 0.5) * width) - for i, v in enumerate(values)) / total - - def _histogram_means(self) -> dict: - """name -> mean of every [histograms] entry, from the post-processing - histograms (the final event sample) when they were filled, otherwise - from the weighted integration histograms.""" - if self.event_histograms is not None: - data = json.loads(self.event_histograms.to_json(self.systematics)) - hists = [(h["name"], h["min"], h["max"], h["bin_values"]) - for h in data] - else: - hists = [(h.name, h.min, h.max, h.bin_values) - for h in self.event_generator.histograms()] - means = {} - for name, hist_min, hist_max, values in hists: - mean = self._histogram_mean(hist_min, hist_max, values) - if mean is not None: - means[name] = mean - return means - - def get_result(self) -> dict: - """Return the run result: cross-section (pb) with MC error, the number - of (unweighted) events, and the mean of every observable declared in the - [histograms] section. Used to build the scan summary.""" - status = self.event_generator.status() - result = {'cross(pb)': status.mean, 'error(pb)': status.error, - 'nb_event': status.count_unweighted} - try: - if self.systematics is not None and self.systematics.weight_count: - summary = json.loads(self.systematics.summary()) - scale = systematics_summary.scale_percentages(summary) - if scale is not None: - result['scale_up(%)'], result['scale_down(%)'] = scale - pdf = systematics_summary.pdf_percentages(summary) - if pdf: - # one column pair; the nominal set's entry comes first - _entry, up, down = pdf[0] - result['pdf_up(%)'], result['pdf_down(%)'] = up, down - except Exception as err: - logger.warning("could not extract the systematics summary: %s", err) - try: - for name, mean in self._histogram_means().items(): - result['<%s>' % name] = mean - except Exception as err: - logger.warning("could not extract observable means: %s", err) - return result - - def _beam_info(self): - """Return (beam_pdg_ids, beam_energies) for the LHE block. - - `incoming` holds the *partonic* initial state (e.g. gluons), not the - 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] - data = self.subprocess_data[0] - incoming = data["incoming"] - flavor0 = data["flavors"][0]["options"][0] - init_pdgs = flavor0[:len(incoming)] - beam_pdgs = [pdg if abs(code) in (81, 82) else code - for code, pdg in zip(incoming, init_pdgs)] - return beam_pdgs, [half_e, half_e] - - def _lhapdf_id(self): - """Central LHAPDF id of the beam PDF set (read from its .info SetIndex), - or -1 for a leptonic beam (no PDF).""" - if self.leptonic or not self.pdf_dir: - return -1 - pdf_set = self.run_card["beam"]["pdf"] - info = os.path.join(self.pdf_dir, pdf_set, "%s.info" % pdf_set) - try: - for line in open(info): - if line.strip().startswith("SetIndex:"): - return int(line.split(":", 1)[1].strip()) - except Exception as err: - logger.warning("could not read LHAPDF id from %s: %s", info, err) - return -1 - - def build_lhe_meta(self, systematics=None): - """Build the LHE header/ metadata: the param_card () and the - run_card.toml () headers plus the beam/PDF/cross-section info - needed by downstream tools (systematics, MadSpin, ...). - - `combine_to_lhe` injects the header itself, so callers that - go through it should leave `systematics` unset; pass it only when - building meta for a writer that bypasses that injection (e.g. - write_lhe_header).""" - beam_pdgs, energies = self._beam_info() - lhaid = self._lhapdf_id() - pdf_group = -1 if self.leptonic else 0 - status = self.event_generator.status() - xsec, err = status.mean, status.error - weight_mode, max_weight = lhe_weight_info(status) - with open(self.param_card_path) as f: - param_text = f.read() - with open(os.path.join("Cards", "run_card.toml")) as f: - run_text = f.read() - headers = [] - # generation commands (model + process): read by MadSpin/reweight/... - proc_card = os.path.join("Cards", "proc_card_mg5.dat") - if os.path.exists(proc_card): - with open(proc_card) as f: - headers.append(ms.LHEHeader(name="MG5ProcCard", content=f.read())) - headers.append(ms.LHEHeader(name="slha", content=param_text)) - headers.append(ms.LHEHeader(name="MG7RunCard", content=run_text)) - # The resolved seed (even when the run_card requested a random one via - # seed = -1), so the run can be reproduced from the LHE file alone. - headers.append(ms.LHEHeader(name="MG7Seed", content=str(self.run_seed))) - if systematics is not None and systematics.weight_ids: - headers.append(ms.LHEHeader( - name="initrwgt", content=systematics.initrwgt(), escape_content=False - )) - return ms.LHEMeta( - beam1_pdg_id=beam_pdgs[0], beam2_pdg_id=beam_pdgs[1], - beam1_energy=energies[0], beam2_energy=energies[1], - beam1_pdf_authors=pdf_group, beam2_pdf_authors=pdf_group, - beam1_pdf_id=lhaid, beam2_pdf_id=lhaid, - weight_mode=weight_mode, - # positional: the pybind arg name for max_weight is non-kwarg-safe - processes=[ms.LHEProcess(xsec, err, max_weight, 1)], - headers=headers, - ) - - def write_lhe_header(self, systematics) -> None: - """Write header.lhe next to events.npy: the
/ blocks - (run card, param card, beam/PDF, cross section) that the npy formats - otherwise drop, with no events.""" - header_path = os.path.join(self.run_path, "header.lhe") - writer = ms.LHEFileWriter(header_path, self.build_lhe_meta(systematics)) - del writer # closes the file (writes the closing tag) - - def build_lhe_completer(self): - all_mcdata = ( - [subproc.build_multi_channel_data() for subproc in self.subprocesses] - if self.merged_subprocess_data is None else - [build_multi_channel_data(meta, self) for meta in self.subprocess_data] - ) - subproc_args = [ - ms.SubprocArgs( - topologies = [topo[0] for topo in mcdata.topologies], - permutations = mcdata.permutations, - diagram_indices = mcdata.diagram_indices, - diagram_color_indices = mcdata.diagram_color_indices, - diagram_propagator_pdgs = mcdata.diagram_propagator_pdgs, - color_flows = meta["color_flows"], - pdg_color_types = { - int(key): value - for key, value in meta["pdg_color_types"].items() - }, - helicities = lhe_helicities(meta, self.run_card), - pdg_ids = [flavor["options"] for flavor in meta["flavors"]], - ) - for mcdata, meta in zip(all_mcdata, self.subprocess_data) - ] - return ms.LHECompleter( - subproc_args=subproc_args, - bw_cutoff=self.run_card["phasespace"]["bw_cutoff"] - ) - - def save_gridpack(self) -> None: - if not self.run_card["gridpack"]["save_gridpack"]: - return - - gridpack_path = os.path.join(self.run_path, "gridpack") - data_path = os.path.join(gridpack_path, "data") - events_path = os.path.join(gridpack_path, "Events") - os.mkdir(gridpack_path) - os.mkdir(data_path) - os.mkdir(events_path) - self.contexts[0].save_globals(os.path.join(data_path, "globals.tar")) - - # one json object per channel, joined by hand so that the channel - # json is not parsed and dumped again - channel_entries = [ - f"{json.dumps(channel.status().name)}:{channel.to_json(self.gridpack_run_mode == "fix_max_weight")}" - for channel in self.event_generator.channels() - ] - with open(os.path.join(data_path, "channels.json"), "w") as f: - f.write("{" + ",".join(channel_entries) + "}") - - lib_path = os.path.join(gridpack_path, "lib") - if self.run_card["gridpack"]["include_source"]: - os.mkdir(lib_path) - shutil.copytree("src", os.path.join(gridpack_path, "src")) - shutil.copytree("SubProcesses", os.path.join(gridpack_path, "SubProcesses")) - shutil.copytree("backend", os.path.join(gridpack_path, "backend")) - else: - shutil.copytree("lib", lib_path) - for _, path in self.custom_function_paths: - target = path if os.path.isabs(path) else os.path.join(gridpack_path, path) - if not os.path.exists(target): - os.makedirs(os.path.dirname(target), exist_ok=True) - shutil.copy(path, target) - - if self.run_card["gridpack"]["include_madspace_source"]: - shutil.copytree( - _MADSPACE_DIR, - os.path.join(gridpack_path, "madspace"), - ignore=shutil.ignore_patterns("build", "install"), - ) - - if self.run_card["gridpack"]["include_madspace"]: - shutil.copytree( - _INSTALL_DIR / "madspace", - os.path.join(gridpack_path, "madspace", "install", "madspace"), - ignore=shutil.ignore_patterns("include"), - ) - - matrix_elements = [] - for subproc in self.subprocess_data: - me_path = subproc["me_path"] - matrix_elements.append(me_path) - - cards_path = os.path.join(gridpack_path, "Cards") - os.mkdir(cards_path) - shutil.copy(os.path.join("Cards", "param_card.dat"), cards_path) - # Full run card with a header noting it is read-only in gridpack context. - import io as _io - _buf = _io.StringIO() - self.run_card.write(_buf) - _header = ( - "# This is the run card used to generate this gridpack.\n" - "# Modifying this file will have no effect on gridpack execution.\n" - "# To change event-generation settings, edit grid_run_card.toml.\n\n" - ) - with open(os.path.join(cards_path, "run_card.toml"), 'w') as _f: - _f.write(_header + _buf.getvalue()) - # Minimal card containing only the settings used by generate_events. - # The gridpack ships the libraries that were actually built, and their - # names carry the resolved backend, so the card it runs from has to name - # that backend as well: a cpu_mode of 'auto' would send the gridpack - # looking for a ..._auto.so that was never built. The full run_card.toml - # written just above keeps 'auto', since that is what was asked for. - device_names = self.run_card["run"]["device"] - if not isinstance(device_names, list): - device_names = [device_names] - resolved_cpu_mode = None - for name, backend in zip(device_names, getattr(self, "backends", [])): - if name.split(":")[0] == "cpu": - resolved_cpu_mode = backend - break - previous_cpu_mode = self.run_card["run"]["cpu_mode"] - if resolved_cpu_mode is not None: - self.run_card["run"]["cpu_mode"] = resolved_cpu_mode - # A gridpack is run to feed a shower or detector chain, which reads - # LHE, so it writes LHE by default whatever format this run used - # (bin/generate_events --output_format still selects an npy output). - previous_output_format = self.run_card["run"]["output_format"] - self.run_card["run"]["output_format"] = "lhe" - # the directories of this run mean nothing on the machine running the gridpack - previous_dirs = (self.run_card["gridpack"]["temp_output_dir"], - self.run_card["gridpack"]["output_dir"]) - self.run_card["gridpack"]["temp_output_dir"] = "" - self.run_card["gridpack"]["output_dir"] = "" - try: - self.run_card.write_gridpack_card( - os.path.join(cards_path, "grid_run_card.toml")) - finally: - self.run_card["run"]["cpu_mode"] = previous_cpu_mode - self.run_card["run"]["output_format"] = previous_output_format - (self.run_card["gridpack"]["temp_output_dir"], - self.run_card["gridpack"]["output_dir"]) = previous_dirs - - bin_path = os.path.join(gridpack_path, "bin") - os.mkdir(bin_path) - gen_events_file = os.path.join(bin_path, "generate_events") - template_dir = os.path.dirname(__file__) - shutil.copy(os.path.join(template_dir, "gridpack.py"), gen_events_file) - shutil.copy(os.path.join(template_dir, "custom_function_loader.py"), bin_path) - os.chmod(gen_events_file, 0o755) - shutil.copy(os.path.join(template_dir, "gridpack_setup.py"), bin_path) - # streaming needs the fixed maximum weights of this mode - if self.gridpack_run_mode == "fix_max_weight": - shutil.copy(os.path.join(template_dir, "event_stream.py"), bin_path) - - data = { - "matrix_elements": matrix_elements, - # loaded before the matrix elements, as in the run - "custom_functions": self.custom_function_paths, - "me_parameters": me_parameters(self.run_card), - "source_hash": ms.SOURCE_HASH, - } - with open(os.path.join(data_path, "data.json"), "w") as f: - json.dump(data, f) - - if self.lhe_completer is None: - self.lhe_completer = self.build_lhe_completer() - self.lhe_completer.save(os.path.join(data_path, "lhe.json")) - self.save_gridpack_lhe_meta(os.path.join(data_path, "lhe_meta.json")) - if self.systematics_data is not None: - systematics_data = dict(self.systematics_data) - # A decay takes alpha_s from a constant grid written into the run - # directory rather than from an LHAPDF set, so the gridpack cannot - # look that file up: ship it and point at the shipped copy. - if self.fixed_alphas_info and \ - systematics_data["nominal_info_file"] == self.fixed_alphas_info: - name = os.path.basename(self.fixed_alphas_info) - shutil.copy(self.fixed_alphas_info, os.path.join(data_path, name)) - systematics_data["nominal_info_file"] = os.path.join("data", name) - with open(os.path.join(data_path, "systematics.json"), "w") as f: - json.dump(systematics_data, f) - if self.run_card["gridpack"]["compress"]: - shutil.make_archive(gridpack_path, "gztar", self.run_path, "gridpack") - shutil.rmtree(gridpack_path) - - def save_gridpack_lhe_meta(self, path) -> None: - """The run-independent part of the LHE header (cards, beams, PDF), for - the gridpack to write complete LHE headers: it adds the cross section - and seed of its own run, which is why those two are left out here.""" - meta = self.build_lhe_meta() - data = {key: getattr(meta, key) for key in ( - "beam1_pdg_id", "beam2_pdg_id", "beam1_energy", "beam2_energy", - "beam1_pdf_authors", "beam2_pdf_authors", "beam1_pdf_id", - "beam2_pdf_id", "weight_mode")} - data["headers"] = [ - {"name": h.name, "content": h.content, - "escape_content": h.escape_content} - for h in meta.headers if h.name != "MG7Seed"] - with open(path, "w") as f: - json.dump(data, f) - - def get_mass(self, pid: int) -> float: - return self.param_card.get_value("mass", pid) - - 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 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 - - -# observables that tell a particle from its antiparticle, and so need the signed -# pdg ids of the actual flavours rather than clean_pids' unsigned representatives -_SIGNED_OBSERVABLES = {"sfos_pair_mass"} - - -def _flavor_options(meta: dict, unmerged_meta: dict | None = None) -> list[list[int]]: - """Every flavour assignment (signed pdg ids, incoming then outgoing) a - subprocess evaluates with the same momenta.""" - if "flavors" in meta: - return [ - option for flavor in meta["flavors"] for option in flavor["options"] - ] - return [ - option - for subproc in meta.get("subprocesses", []) - for option in _flavor_options(unmerged_meta[subproc]) - ] - - -def observable_pids( - meta: dict, observable_kwargs: dict, unmerged_meta: dict | None = None -) -> list[int]: - """The pdg ids an observable of this subprocess is built with. - - Most observables only select particles by group, which the unsigned - representatives of clean_pids do. One that pairs a particle with its - antiparticle needs the real signed flavours: those of the first flavour - assignment, provided every assignment of the subprocess forms the same - pairs -- the cut is applied once to momenta that all of them share. - """ - all_pids = clean_pids(meta["incoming"]) + clean_pids(meta["outgoing"]) - if observable_kwargs.get("observable") not in _SIGNED_OBSERVABLES: - return all_pids - options = _flavor_options(meta, unmerged_meta) - if not options: - raise ValueError( - f"{observable_kwargs.get('name')}: no flavour information for " - f"subprocess {meta['incoming']} > {meta['outgoing']}" - ) - n_in = len(meta["incoming"]) - groups = [set(group) for group in observable_kwargs["select_pids"]] - if len(groups) == 1: - groups = groups * 2 - - def sfos_pairs(pids): - out = range(n_in, len(pids)) - return frozenset( - frozenset((i, j)) - for i in out if pids[i] in groups[0] - for j in out if pids[j] in groups[1] - if i != j and pids[i] != 0 and pids[i] == -pids[j] - ) - - reference = sfos_pairs(options[0]) - if any(sfos_pairs(option) != reference for option in options[1:]): - raise ValueError( - f"{observable_kwargs.get('name')}: the flavours grouped into subprocess " - f"{meta['incoming']} > {meta['outgoing']} pair the particles " - "differently, so a same-flavour opposite-sign cut cannot be applied " - "to their common momenta; regenerate the process after " - "'set apply_flavor_grouping False'" - ) - return list(options[0]) - - -def build_observable(pids: list[int], n_in: int, observable_kwargs: dict): - """ms.Observable of a subprocess with `n_in` incoming particles. - - madspace takes the first two momenta to be the beams (ignore_incoming), so - for a 1 -> N decay it would drop the first decay product from every - selection. There the selection runs over all the momenta instead, with the - incoming particle set to pdg id 0, which is in no group. - """ - if n_in == 2: - return ms.Observable(pids, **observable_kwargs) - pids = [0] * n_in + list(pids[n_in:]) - return ms.Observable(pids, **dict(observable_kwargs, ignore_incoming=False)) - - -def decay_products(meta: dict, unmerged_meta: dict | None = None) -> set[int]: - """Outgoing positions (0-based, beams excluded) of the particles that come - from an on-shell (decay-chain) propagator, as MadEvent's check_decay marks - them. Every channel of a subprocess shares its decay chains, so the first - channel tells, read through its first diagram's permutation.""" - if not meta["channels"]: - return set() - channel = meta["channels"][0] - topo_channel = channel - if "vertices" not in channel: - topo_channel = unmerged_meta[channel["subprocess"]]["channels"][channel["channel"]] - on_shell = topo_channel.get("on_shell_propagators", []) - if not on_shell: - return set() - vertices = topo_channel["vertices"] - # the closing vertex creates no propagator - created = {vertex[-1]: vertex[:-1] for vertex in vertices[:-1]} - - def legs(name): - if name in created: - return set().union(*(legs(child) for child in created[name])) - return {name} - - n_in = len(meta["incoming"]) - n_out = len(meta["outgoing"]) - products = set() - for prop in on_shell: - external = legs(f"p{prop}") - incoming = {leg for leg in external if leg.startswith("i")} - outgoing = {int(leg[1:]) for leg in external if leg.startswith("o")} - if not incoming: - products |= outgoing - elif len(incoming) == n_in: - # built from the beams: the decay is everything else - products |= set(range(n_out)) - outgoing - # event[i] = topology[permutation[i]] - permutation = channel["diagrams"][0]["permutation"] - return { - i - n_in for i in range(n_in, len(permutation)) - if permutation[i] - n_in in products - } - - -def pid_is_qcd(pid: int): - return abs(pid) in [21, 1, 2, 3, 4, 5, 6, 81] - - -def build_topologies( - incoming_masses: list[float], - outgoing_masses: list[float], - channel: dict, - process: MadgraphProcess -) -> list[ms.Topology]: - propagators = [] - for i, (pid, signed_pid) in enumerate(zip( - clean_pids(channel["propagators"]), channel["propagators"] - )): - mass = process.get_mass(pid) - width = process.get_width(pid) - bw_cutoff = process.run_card["phasespace"]["bw_cutoff"] - if i in channel["on_shell_propagators"]: - e_min = mass - bw_cutoff * width - e_max = mass + bw_cutoff * width - else: - e_min = 0 - e_max = 0 - # $-excluded propagator: the matrix element vanishes within bw_cutoff - # widths of the pole, so do not pile the channel's points up there - # (a flat density in that window instead of the Breit-Wigner peak). - # Channels written before the exporter recorded them have none. - flat_window = 0. - if i in channel.get("dollar_propagators", ()) and width > 0: - flat_window = float(bw_cutoff) - propagators.append(ms.Propagator( - mass=mass, - width=width, - integration_order=0, - e_min=e_min, - e_max=e_max, - pdg_id=signed_pid, - flat_window=flat_window, - )) - vertices = channel["vertices"] - diag = ms.Diagram( - incoming_masses, outgoing_masses, propagators, vertices - ) - return ms.Topology.topologies(diag) - - -def build_multi_channel_data( - meta: dict, process: MadgraphProcess, unmerged_meta: dict | None = None -) -> MultiChannelData: - incoming_masses = [ - process.get_mass(pid) for pid in clean_pids(meta["incoming"]) - ] - outgoing_masses = [ - process.get_mass(pid) for pid in clean_pids(meta["outgoing"]) - ] - - if unmerged_meta is None: - diagram_count = meta["diagram_count"] - amp2_remaps = [[-1] * diagram_count] - else: - amp2_remaps = [ - [-1] * unmerged_meta[subproc]["diagram_count"] - for subproc in meta["subprocesses"] - ] - symfact = [] - topologies = [] - permutations = [] - channel_indices = [] - channel_weight_indices = [] - diagram_indices = [] - diagram_color_indices = [] - diagram_propagator_pdgs = [] - active_flavors = [] - channel_index = 0 - qcd_s_channel_count = [] - - for channel in meta["channels"]: - if unmerged_meta is None: - topo_channel = channel - else: - topo_subproc = channel["subprocess"] - topo_channel_index = channel["channel"] - topo_channel = unmerged_meta[topo_subproc]["channels"][topo_channel_index] - chan_topologies = build_topologies( - incoming_masses, outgoing_masses, topo_channel, process - ) - topo_count = len(chan_topologies) - if topo_count == 0: - continue - - topo = chan_topologies[0] - s_chan_count = [0] * len(topo.decays) - non_qcd = [False] * len(topo.decays) - pdg_ids = [decay.pdg_id for decay in topo.decays] - for i, pid in zip(topo.outgoing_indices, meta["outgoing"]): - pdg_ids[i] = pid - for decay in reversed(topo.decays): - if len(decay.child_indices) == 0: - continue - if decay.index == 0 and topo.t_propagator_count > 0: - s_chan_count[0] = sum(s_chan_count[i] for i in decay.child_indices) - break - - is_qcd = pid_is_qcd(pdg_ids[decay.index]) and all( - pid_is_qcd(pdg_ids[i]) for i in decay.child_indices - ) - is_non_qcd = decay.on_shell or not is_qcd or any( - non_qcd[i] for i in decay.child_indices - ) - non_qcd[decay.index] = is_non_qcd - s_chan_count[decay.index] = ( - 0 if is_non_qcd else sum(s_chan_count[i] for i in decay.child_indices) + 1 - ) - qcd_s_channel_count.append(s_chan_count[0]) - - diagrams = channel["diagrams"] - chan_permutations = [d["permutation"] for d in diagrams] - if unmerged_meta is None: - amp2_remaps[0][diagrams[0]["diagram"]] = channel_index - else: - for amp2_remap, diag in zip(amp2_remaps, diagrams[0]["diagram"]): - if diag != -1: - amp2_remap[diag] = channel_index - - channel_index_first = channel_index - symfact_index_first = len(symfact) - channel_index += 1 - symfact.extend([None] * topo_count) - for d in diagrams[1:]: - if unmerged_meta is None: - amp2_remaps[0][d["diagram"]] = channel_index - else: - for amp2_remap, diag in zip(amp2_remaps, d["diagram"]): - if diag != -1: - amp2_remap[diag] = channel_index - channel_index += 1 - symfact.extend(range(symfact_index_first, symfact_index_first + topo_count)) - - topologies.append(chan_topologies) - permutations.append(chan_permutations) - channel_indices.append(list(range(channel_index_first, channel_index))) - channel_weight_indices.append([ - [ - symfact_index_first + topo_index + i * topo_count - for i in range(len(chan_permutations)) - ] - for topo_index in range(topo_count) - ]) - diagram_indices.append([d["diagram"] for d in diagrams]) - if unmerged_meta is None: - diagram_color_indices.append([d["active_colors"] for d in diagrams]) - diagram_propagator_pdgs.append( - [d["propagator_pdgs"] for d in diagrams] - ) - active_flavors.append([d["active_flavors"] for d in diagrams]) - - return MultiChannelData( - amp2_remaps, - symfact, - topologies, - permutations, - channel_indices, - channel_weight_indices, - diagram_indices, - diagram_color_indices, - diagram_propagator_pdgs, - active_flavors, - qcd_s_channel_count, - ) - - -class MadgraphSubprocess: - def __init__( - self, - process: MadgraphProcess, - meta: dict, - subproc_id: int, - unmerged_meta: dict | None = None - ): - self.process = process - self.meta = meta - self.subproc_id = subproc_id - self.multi_channel_data = None - - self.unmerged_meta = None - if unmerged_meta is None: - api_path_formats = [self.meta["me_path"]] - member_metas = [self.meta] - else: - api_path_formats = [] - member_metas = [] - for subproc in self.meta["subprocesses"]: - submeta = unmerged_meta[subproc] - api_path_formats.append(submeta["me_path"]) - member_metas.append(submeta) - if len(api_path_formats) == 1: - self.meta = submeta - else: - self.unmerged_meta = unmerged_meta - - # The libraries were all built up front by MadgraphProcess.compile_matrix_elements - all_api_paths = [ - [api_path_format.format(device=backend) for backend in self.process.backends] - for api_path_format in api_path_formats - ] - - self.incoming_masses = [ - self.process.get_mass(pid) for pid in clean_pids(self.meta["incoming"]) - ] - self.outgoing_masses = [ - self.process.get_mass(pid) for pid in clean_pids(self.meta["outgoing"]) - ] - self.particle_count = len(self.incoming_masses) + len(self.outgoing_masses) - # MadEvent's cut_decays = F: the decay products of on-shell propagators - # are left out of every cut selection (a pair cut then needs both - # members to be cut), while the histograms still see them - if self.process.run_card["phasespace"]["cut_decays"]: - uncut = set() - else: - uncut = decay_products(self.meta, self.unmerged_meta) - n_in = len(self.incoming_masses) - - def cut_pids(observable_kwargs): - pids = observable_pids(self.meta, observable_kwargs, self.unmerged_meta) - # pdg id 0 is in no group, so the particle is never selected - return [0 if i - n_in in uncut else pid for i, pid in enumerate(pids)] - - # custom cuts and scales see the PDG ids of the first flavor, like the - # built-in cuts, which select by flavor-independent groups - representative_pids = [self.flavor_pids()[0]] - custom_cut = self.process.custom_function( - "cut", self.particle_count, representative_pids - ) - if len(self.process.cut_data) > 0: - self.cuts = ms.Cuts([ - ms.CutItem( - observable=build_observable( - cut_pids(cut_item.observable_kwargs), - n_in, cut_item.observable_kwargs - ), - min=cut_item.min, - max=cut_item.max, - mode=cut_item.mode, - ) - for cut_item in self.process.cut_data - ], custom_cut) - elif custom_cut is not None: - self.cuts = ms.Cuts(self.particle_count, custom_cut) - else: - self.cuts = None - # the integration histograms are functions of the momenta; the weight - # distribution is a property of the final event sample, so it is only - # filled by MadgraphProcess.build_event_histograms at combine time - # They cost an observable evaluation per integration point, and the - # plots/HwU are drawn from the post-processing ones instead, so they - # are only filled on request ([run] weighted_histograms). - momentum_hists = [item for item in self.process.hist_data - if not item.from_weight] \ - if self.process.run_card["run"]["weighted_histograms"] else [] - self.histograms = ( - ms.ObservableHistograms([ - ms.HistItem( - observable=build_observable( - observable_pids( - self.meta, hist_item.observable_kwargs, self.unmerged_meta - ), - n_in, hist_item.observable_kwargs - ), - min=hist_item.min, - max=hist_item.max, - bin_count=hist_item.bin_count, - ) - for hist_item in momentum_hists - ]) - if momentum_hists - else None - ) - - self.scale = ms.EnergyScale( - particle_count=self.particle_count, - custom_scale=self.process.custom_function( - "scale", self.particle_count, representative_pids - ), - **self.process.scale_kwargs - ) - - if self.process.run_card["run"]["dummy_matrix_element"]: - self.matrix_elements = [None] * len(all_api_paths) - else: - self.matrix_elements = [] - for api_paths in all_api_paths: - for context, api_path in zip(self.process.contexts, api_paths): - mat = context.load_matrix_element( - api_path, self.process.param_card_path, - me_parameters(self.process.run_card) - ) - self.matrix_elements.append(mat) - self.surrogates = [] - if self.process.surrogate is not None: - for path_format, member_meta in zip(api_path_formats, member_metas): - self.surrogates.append(self.load_surrogate(path_format, member_meta)) - - def load_surrogate(self, me_path_format: str, member_meta: dict): - """The [surrogate] of one matrix element of the subprocess.""" - name = re.sub( - r"^libmadmatrix_(.*)_\{device\}\.so$", r"\1", os.path.basename(me_path_format) - ) - # PDG ids for each flavor index of the matrix element - flavors = member_meta["flavors"] - pdg_ids = [flavors[0]["options"][0]] * (max(f["index"] for f in flavors) + 1) - for flavor in flavors: - pdg_ids[flavor["index"]] = flavor["options"][0] - settings = self.process.surrogate - apis = load_surrogate( - ms, - self.process.contexts, - self.process.backends, - self.process.device_types, - settings["function"], - name, - pdg_ids, - settings["param_card"], - ) - if len({api.index() for api in apis}) != 1: - raise RuntimeError("surrogate loaded at different indices") - return apis[0] - - def flavor_pids(self) -> list[list[int]]: - """PDG ids of the external particles for every flavor of the - integrand, in the order of its flavor index.""" - pids = [] - for flav in self.meta["flavors"]: - if self.unmerged_meta is not None: - subproc_index = self.meta["subprocesses"][flav["subprocess"]] - flav = self.unmerged_meta[subproc_index]["flavors"][flav["flavor"]] - pids.append(flav["options"][0]) - return pids - - def build_multi_channel_data(self) -> MultiChannelData: - if self.multi_channel_data is not None: - return self.multi_channel_data - self.multi_channel_data = build_multi_channel_data( - self.meta, self.process, self.unmerged_meta - ) - return self.multi_channel_data - - def build_multichannel_phasespace(self) -> PhaseSpace: - mcdata = self.build_multi_channel_data() - channel_count = sum(len(topos) for topos in mcdata.topologies) - drop_threshold = self.process.run_card["phasespace"]["drop_qcd_s_channel"] - if drop_threshold >= 0 and channel_count > drop_threshold: - mcdata = self.drop_qcd_s_channels(mcdata) - mcdata = self.drop_empty_channels(mcdata) - - channels = [] - t_channel_mode = self.t_channel_mode( - self.process.run_card["phasespace"]["t_channel"] - ) - for channel_id, (chan_topologies, chan_permutations, chan_indices, active_flavors) in enumerate(zip( - mcdata.topologies, mcdata.permutations, mcdata.channel_weight_indices, - mcdata.active_flavors - )): - topo_count = len(chan_topologies) - for topo_index, (topo, indices) in enumerate(zip(chan_topologies, chan_indices)): - mapping = ms.PhaseSpaceMapping( - chan_topologies[0], - self.process.e_cm, - t_channel_mode=t_channel_mode, - cuts=self.cuts, - invariant_power=self.process.run_card["phasespace"]["invariant_power"], - permutations=chan_permutations, - leptonic=self.process.leptonic, - ) - prefix = f"subproc{self.subproc_id}.channel{channel_id}" - if topo_count > 1: - prefix += f".subchan{topo_index}" - discrete_sym, discrete_flavor = self.build_discrete( - len(chan_permutations), len(self.meta["flavors"]), prefix - ) - channels.append(Channel( - phasespace_mapping = mapping, - adaptive_mapping = self.build_vegas(mapping, prefix), - discrete_sym = discrete_sym, - discrete_flavor = discrete_flavor, - channel_weight_indices = indices, - name = f"{channel_id}", - active_flavors = active_flavors, - )) - - remapped_chan_count = sum( - len(indices) for indices in mcdata.channel_indices - ) - if self.process.run_card["phasespace"]["sde_strategy"] == "denominators": - prop_chan_weights = ms.PropagatorChannelWeights( - [topo[0] for topo in mcdata.topologies], mcdata.permutations, - mcdata.channel_indices - ) - chan_weight_remap = [] - else: - prop_chan_weights = None - chan_weight_remap = [ - [ - len(mcdata.symfact) if remap == -1 else remap - for remap in amp2_remap - ] - for amp2_remap in mcdata.amp2_remaps - ] - - if any(len(topos) > 1 for topos in mcdata.topologies): - subchan_weights = ms.SubchannelWeights( - mcdata.topologies, mcdata.permutations, mcdata.channel_indices - ) - else: - subchan_weights = None - - return PhaseSpace( - mode="multichannel", - channels=channels, - first_chan_weight_remap=chan_weight_remap, - first_remapped_chan_count=remapped_chan_count, - symfact=mcdata.symfact, - prop_chan_weights=prop_chan_weights, - subchan_weights=subchan_weights, - ) - - def build_flat_phasespace(self) -> PhaseSpace: - mapping = ms.PhaseSpaceMapping( - self.incoming_masses + self.outgoing_masses, - self.process.e_cm, - mode=self.t_channel_mode(self.process.run_card["phasespace"]["flat_mode"]), - cuts=self.cuts, - leptonic=self.process.leptonic, - ) - prefix = f"subproc{self.subproc_id}.flat" - discrete_sym, discrete_flavor = self.build_discrete( - 1, len(self.meta["flavors"]), prefix - ) - channel = Channel( - phasespace_mapping = mapping, - adaptive_mapping = self.build_vegas(mapping, prefix), - discrete_sym = discrete_sym, - discrete_flavor = discrete_flavor, - channel_weight_indices = [0], - name = "F", - active_flavors = [], - ) - if self.unmerged_meta is None: - remap = [list(range(self.meta["diagram_count"]))] - else: - remap = [ - list(range(self.unmerged_meta[subproc]["diagram_count"])) - for subproc in self.meta["subprocesses"] - ] - return PhaseSpace( - mode="flat", - channels=[channel], - first_chan_weight_remap=remap, - first_remapped_chan_count=1, - symfact=[None], - ) - - def simplify_phasespace( - self, - multi_phasespace: PhaseSpace, - cross_sections: list[float] - ) -> PhaseSpace | None: - assert multi_phasespace.mode == "multichannel" - - threshold = 1 - self.process.run_card["phasespace"]["combine_channel_threshold"] - kept_channels = [] - tot_cs = sum(cross_sections) - cum_cs = 0. - seen_active_flavors = set() - #seen_resonances = set() - for index, (cs, chan) in sorted( - enumerate(zip(cross_sections, multi_phasespace.channels)), - key=lambda pair: pair[1][0], - reverse=True - ): - cum_cs += cs - has_unseen_flavors = False - has_unseen_resonances = False - for flavs in chan.active_flavors: - for flav in flavs: - if flav not in seen_active_flavors: - has_unseen_flavors = True - seen_active_flavors.add(flav) - #for resonance in chan.resonances: - # if resonance not in seen_resonances: - # has_unseen_resonances = True - # seen_resonances.add(flav) - if has_unseen_flavors or has_unseen_resonances or not tot_cs \ - or cum_cs / tot_cs < threshold: - kept_channels.append(index) - if len(kept_channels) >= len(cross_sections) - 1: - return multi_phasespace - - channels = [] - channel_map = {} - symfact = [] - for old_chan_index in kept_channels: - channel = multi_phasespace.channels[old_chan_index] - perm_count = max(1, channel.phasespace_mapping.channel_count()) - channel_index = len(symfact) - symfact.append(None) - symfact.extend([channel_index] * (perm_count - 1)) - channel_map.update({ - old_index: new_index - for new_index, old_index in enumerate( - channel.channel_weight_indices, start=channel_index - ) - }) - channels.append(Channel( - phasespace_mapping = channel.phasespace_mapping, - adaptive_mapping = channel.adaptive_mapping, - discrete_sym = channel.discrete_sym, - discrete_flavor = channel.discrete_flavor, - channel_weight_indices = list(range( - channel_index, channel_index + perm_count - )), - name = channel.name, - active_flavors = channel.active_flavors, - event_generator = channel.event_generator, - )) - - flat_phasespace = self.build_flat_phasespace() - flat_channel = flat_phasespace.channels[0] - channels.append(Channel( - phasespace_mapping = flat_channel.phasespace_mapping, - adaptive_mapping = flat_channel.adaptive_mapping, - discrete_sym = flat_channel.discrete_sym, - discrete_flavor = flat_channel.discrete_flavor, - channel_weight_indices = [len(symfact)], - name = flat_channel.name, - active_flavors = flat_channel.active_flavors, - )) - flat_index = len(symfact) - symfact.append(None) - channel_map[len(multi_phasespace.symfact)] = len(symfact) - if multi_phasespace.subchan_weights is None and len(multi_phasespace.first_chan_weight_remap) > 0: - first_chan_weight_remap = [ - [ - channel_map.get(remap, flat_index) - for remap in cw_remap - ] - for cw_remap in multi_phasespace.first_chan_weight_remap - ] - first_remapped_chan_count = len(symfact) - second_chan_weight_remap = [] - second_remapped_chan_count = 0 - else: - first_chan_weight_remap = multi_phasespace.first_chan_weight_remap - first_remapped_chan_count = multi_phasespace.first_remapped_chan_count - chan_count = multi_phasespace.first_remapped_chan_count if multi_phasespace.subchan_weights is None else multi_phasespace.subchan_weights.channel_count() - second_chan_weight_remap = [ - channel_map.get(i, flat_index) - for i in range(chan_count) - ] - second_remapped_chan_count = len(symfact) - - return PhaseSpace( - mode="both", - channels=channels, - first_chan_weight_remap=first_chan_weight_remap, - first_remapped_chan_count=first_remapped_chan_count, - second_chan_weight_remap=second_chan_weight_remap, - second_remapped_chan_count=second_remapped_chan_count, - symfact=symfact, - prop_chan_weights=multi_phasespace.prop_chan_weights, - subchan_weights=multi_phasespace.subchan_weights, - ) - - def drop_qcd_s_channels(self, mcdata: MultiChannelData) -> MultiChannelData: - """Drop channels with non-resonant QCD s-channel propagators, to reduce the - channel count for processes with many diagrams. Channels are grouped by the - number of QCD s-channels, allowing for more if necessary to map out all flavor - indices. Channel weights belonging to a dropped channel are left unmapped.""" - groups_by_s_count = {} - for index, count in enumerate(mcdata.qcd_s_channel_count): - groups_by_s_count.setdefault(count, []).append(index) - - covered_flavors = set() - kept_groups = set() - for s_count in sorted(groups_by_s_count): - s_count_groups = groups_by_s_count[s_count] - if s_count == 0: - selected = s_count_groups - else: - selected = [ - index - for index in s_count_groups - if any( - flav not in covered_flavors - for flavs in mcdata.active_flavors[index] - for flav in flavs - ) - ] - kept_groups.update(selected) - for index in selected: - for flavs in mcdata.active_flavors[index]: - covered_flavors.update(flavs) - - return self.keep_channel_groups(mcdata, sorted(kept_groups)) - - def drop_empty_channels(self, mcdata: MultiChannelData) -> MultiChannelData: - """Drop the channels whose phase space the cuts exclude entirely: a pair - mass cut whose floor lies at or above the top of an on-shell window of - the channel (ms.PhaseSpaceMapping.empty). Every point such a channel - generates fails the cuts, so it could only ever fail to find one. - - If that leaves a flavour with no channel at all, its cross section is - zero -- typically a cut that excludes the window of a decay chain every - diagram shares -- and the run cannot sample it: that is reported as an - error naming the cause, rather than as the survey finding no point - passing the cuts.""" - if self.cuts is None: - return mcdata - t_channel_mode = self.t_channel_mode( - self.process.run_card["phasespace"]["t_channel"] - ) - empty = [ - ms.PhaseSpaceMapping( - chan_topologies[0], - self.process.e_cm, - t_channel_mode=t_channel_mode, - cuts=self.cuts, - invariant_power=self.process.run_card["phasespace"]["invariant_power"], - permutations=chan_permutations, - leptonic=self.process.leptonic, - ).empty() - for chan_topologies, chan_permutations in zip( - mcdata.topologies, mcdata.permutations - ) - ] - if not any(empty): - return mcdata - kept_groups = [index for index, is_empty in enumerate(empty) if not is_empty] - covered_flavors = { - flav - for index in kept_groups - for flavs in mcdata.active_flavors[index] - for flav in flavs - } - uncovered = sorted({ - flav - for index, is_empty in enumerate(empty) if is_empty - for flavs in mcdata.active_flavors[index] - for flav in flavs - } - covered_flavors) - if uncovered: - raise ValueError( - f"subprocess {self.meta['incoming']} > {self.meta['outgoing']}: " - "the pair mass cuts exclude the whole on-shell window " - f"(bw_cutoff = {self.process.run_card['phasespace']['bw_cutoff']}) " - f"of every channel of flavour(s) {uncovered}, so their cross " - "section is zero; loosen the cut or remove that decay" - ) - logger.info( - "subprocess %d: %d channel(s) dropped, the cuts exclude the " - "on-shell window they sample", self.subproc_id, len(empty) - len(kept_groups) - ) - return self.keep_channel_groups(mcdata, kept_groups) - - def keep_channel_groups( - self, mcdata: MultiChannelData, kept_groups: list[int] - ) -> MultiChannelData: - """Rebuild the multi-channel data with only the channel groups (indices - into mcdata.topologies) in kept_groups. Channel weights belonging to a - dropped channel are left unmapped.""" - if len(kept_groups) == len(mcdata.topologies): - return mcdata - - amp2_remaps = [[-1] * len(remap) for remap in mcdata.amp2_remaps] - symfact = [] - topologies = [] - permutations = [] - channel_indices = [] - channel_weight_indices = [] - diagram_indices = [] - diagram_color_indices = [] - diagram_propagator_pdgs = [] - active_flavors = [] - qcd_s_channel_count = [] - channel_index = 0 - - for group in kept_groups: - chan_topologies = mcdata.topologies[group] - chan_permutations = mcdata.permutations[group] - chan_diagram_indices = mcdata.diagram_indices[group] - topo_count = len(chan_topologies) - - channel_index_first = channel_index - symfact_index_first = len(symfact) - for i, diag in enumerate(chan_diagram_indices): - if self.unmerged_meta is None: - amp2_remaps[0][diag] = channel_index - else: - for amp2_remap, d in zip(amp2_remaps, diag): - if d != -1: - amp2_remap[d] = channel_index - if i == 0: - symfact.extend([None] * topo_count) - else: - symfact.extend(range(symfact_index_first, symfact_index_first + topo_count)) - channel_index += 1 - - topologies.append(chan_topologies) - permutations.append(chan_permutations) - channel_indices.append(list(range(channel_index_first, channel_index))) - channel_weight_indices.append([ - [ - symfact_index_first + topo_index + i * topo_count - for i in range(len(chan_permutations)) - ] - for topo_index in range(topo_count) - ]) - diagram_indices.append(chan_diagram_indices) - if mcdata.diagram_color_indices: - diagram_color_indices.append(mcdata.diagram_color_indices[group]) - if mcdata.diagram_propagator_pdgs: - diagram_propagator_pdgs.append( - mcdata.diagram_propagator_pdgs[group] - ) - active_flavors.append(mcdata.active_flavors[group]) - qcd_s_channel_count.append(mcdata.qcd_s_channel_count[group]) - - return MultiChannelData( - amp2_remaps, - symfact, - topologies, - permutations, - channel_indices, - channel_weight_indices, - diagram_indices, - diagram_color_indices, - diagram_propagator_pdgs, - active_flavors, - qcd_s_channel_count, - ) - - def set_madnis_auto_settings(self, rsd: float): - madnis_args = self.process.run_card["madnis"] - n_out = len(self.meta["outgoing"]) - n_events = self.process.run_card["generation"]["events"] - is_gridpack = self.process.run_card["gridpack"]["save_gridpack"] - train_batches = madnis_args["train_batches"] - hidden_dim = min(max(int((7 * rsd) / 32) * 32 + 64, 64), 256) - flow_layers = 3 if rsd < 32 else 4 - lr = min(max((10000 - train_batches) / 8000 * 7e-4 + 3e-4, 3e-4), 1e-3) - if n_out <= 2: - enable = False - elif n_out == 3: - enable = rsd > 10. or n_events > 1000000 or is_gridpack - else: - enable = True - fixed_cwnet_fraction = max(0.33, 1.0 - 10000. / train_batches) - batch_size_per_channel = min(max(int((7 * rsd) / 32) * 32 + 64, 128), 512) - - self.madnis_settings = { - "enable": enable, - "flow_layers": flow_layers, - "flow_hidden_dim": hidden_dim, - "discrete_hidden_dim": hidden_dim, - "cwnet_hidden_dim": hidden_dim, - "train_batches": train_batches, - "lr": lr, - "fixed_cwnet_fraction": fixed_cwnet_fraction, - "batch_size_per_channel": batch_size_per_channel, - } - for key, value in self.madnis_settings.items(): - run_card_value = madnis_args[key] - if run_card_value != "auto": - self.madnis_settings[key] = run_card_value - - def build_madnis(self, phasespace: PhaseSpace) -> PhaseSpace: - madnis_args = self.process.run_card["madnis"] - # Shared across all networks below: initialize_globals() derives an - # independent, non-colliding stream per tensor from this one base seed. - seed = self.process.run_seed - channels = [] - for channel_id, channel in enumerate(phasespace.channels): - prefix = f"subproc{self.subproc_id}.channel{channel_id}" - cond_dim = 0 - - flow_dim = channel.phasespace_mapping.random_dim() - flow = ms.Flow( - input_dim=flow_dim, - condition_dim=cond_dim, - prefix=prefix, - bin_count=madnis_args["flow_spline_bins"], - subnet_hidden_dim=self.madnis_settings["flow_hidden_dim"], - subnet_layers=self.madnis_settings["flow_layers"], - subnet_activation=self.activation(madnis_args["flow_activation"]), - invert_spline=madnis_args["flow_invert_spline"], - ) - if channel.adaptive_mapping is None: - flow.initialize_globals(self.process.contexts[0], seed) - else: - flow.initialize_from_vegas( - self.process.contexts[0], channel.adaptive_mapping.grid_name(), - seed - ) - cond_dim += flow_dim - - # discrete_sym runs after the adaptive map, so it can condition on its latent. - discrete_sym = channel.discrete_sym - if discrete_sym is not None: - perm_count = channel.phasespace_mapping.channel_count() - discrete_sym = ms.DiscreteFlow( - option_counts=[perm_count], - prefix=f"{prefix}.discrete_flow_sym", - dims_with_prior=[], - condition_dim=cond_dim, - subnet_hidden_dim=self.madnis_settings["discrete_hidden_dim"], - subnet_layers=madnis_args["discrete_layers"], - subnet_activation=self.activation(madnis_args["discrete_activation"]), - ) - discrete_sym.initialize_globals(self.process.contexts[0], seed) - cond_dim += perm_count - - discrete_flavor = channel.discrete_flavor - if discrete_flavor is not None: - discrete_flavor = ms.DiscreteFlow( - option_counts=[len(self.meta["flavors"])], - prefix=f"{prefix}.discrete_flow_flavor", - dims_with_prior=[0], - condition_dim=cond_dim, - subnet_hidden_dim=self.madnis_settings["discrete_hidden_dim"], - subnet_layers=madnis_args["discrete_layers"], - subnet_activation=self.activation(madnis_args["discrete_activation"]), - ) - discrete_flavor.initialize_globals(self.process.contexts[0], seed) - - channels.append(Channel( - phasespace_mapping = channel.phasespace_mapping, - adaptive_mapping = flow, - discrete_sym = discrete_sym, - discrete_flavor = discrete_flavor, - channel_weight_indices = channel.channel_weight_indices, - name = channel.name, - active_flavors = channel.active_flavors, - )) - - return PhaseSpace( - mode="both", - channels=channels, - first_chan_weight_remap=phasespace.first_chan_weight_remap, - first_remapped_chan_count=phasespace.first_remapped_chan_count, - second_chan_weight_remap=phasespace.second_chan_weight_remap, - second_remapped_chan_count=phasespace.second_remapped_chan_count, - symfact=phasespace.symfact, - cwnet=self.build_cwnet(len(phasespace.symfact)), - prop_chan_weights=phasespace.prop_chan_weights, - subchan_weights=phasespace.subchan_weights, - ) - - def build_vegas(self, mapping: ms.PhaseSpaceMapping, prefix: str) -> ms.VegasMapping: - if not self.process.run_card["vegas"]["enable"]: - return None - - vegas = ms.VegasMapping( - mapping.random_dim(), - self.process.run_card["vegas"]["bins"], - prefix, - ) - for context in self.process.contexts: - vegas.initialize_globals(context) - return vegas - - def build_discrete( - self, permutation_count: int, flavor_count: int, prefix: str - ) -> tuple[ms.DiscreteSampler | None, ms.DiscreteSampler | None]: - is_adaptive = self.process.run_card["phasespace"]["adaptive_symmetry_sampling"] - if is_adaptive and permutation_count > 1: - discrete_sym = ms.DiscreteSampler( - [permutation_count], f"{prefix}.discrete_sym" - ) - for context in self.process.contexts: - discrete_sym.initialize_globals(context) - else: - discrete_sym = None - - if flavor_count > 1: - # the flavor probabilities are conditioned on the PDF prior; a - # leptonic process (or a decay) has none, and the Integrand then - # passes no condition: declaring one anyway fails to build - # ("keys and values must have the same size"), e.g. z > q q~ with - # both u- and d-type rows - discrete_flavor = ms.DiscreteSampler( - [flavor_count], f"{prefix}.discrete_flavor", - [] if self.process.leptonic else [0] - ) - for context in self.process.contexts: - discrete_flavor.initialize_globals(context) - else: - discrete_flavor = None - - return discrete_sym, discrete_flavor - - def build_cwnet(self, channel_count: int) -> ms.ChannelWeightNetwork: - #if channel_count == 1: - # return None - madnis_args = self.process.run_card["madnis"] - cwnet = ms.ChannelWeightNetwork( - channel_count=channel_count, - particle_count=self.particle_count, - hidden_dim=self.madnis_settings["cwnet_hidden_dim"], - layers=madnis_args["cwnet_layers"], - activation=self.activation(madnis_args["cwnet_activation"]), - prefix=f"subproc{self.subproc_id}.cwnet", - ) - cwnet.initialize_globals( - self.process.contexts[0], self.process.run_seed - ) - return cwnet - - def t_channel_mode(self, name: str) -> ms.PhaseSpaceMapping.TChannelMode: - modes = { - "propagator": ms.PhaseSpaceMapping.propagator, - "rambo": ms.PhaseSpaceMapping.rambo, - "chili": ms.PhaseSpaceMapping.chili, - } - if name in modes: - return modes[name] - else: - raise ValueError(f"Invalid t-channel mode '{name}'") - - def activation(self, name: str) -> ms.MLP.Activation: - activations = { - "relu": ms.MLP.relu, - "leaky_relu": ms.MLP.leaky_relu, - "elu": ms.MLP.elu, - "gelu": ms.MLP.gelu, - "sigmoid": ms.MLP.sigmoid, - "softplus": ms.MLP.softplus, - } - if name in activations: - return activations[name] - else: - raise ValueError(f"Invalid activation function '{name}'") - - def build_integrands( - self, - phasespace: PhaseSpace, - madnis_training: bool = False, - drop_cuts_and_rescale: bool = False - ) -> list[ms.Integrand]: - flavors = [] - flavor_remap = [] - flavor_factors = [] - flavor_mirror = [] - flavor_diff_xs_indices = [] - flavor_subproc_indices = [] - flavor_per_subproc_remap = [] - - for flav in self.meta["flavors"]: - if self.unmerged_meta is not None: - diff_xs_index = flav["subprocess"] - subproc_index = self.meta["subprocesses"][diff_xs_index] - ps_flavor = flav["flavor"] - flavor_diff_xs_indices.append(diff_xs_index) - flavor_subproc_indices.append(subproc_index) - flavor_per_subproc_remap.append(ps_flavor) - flav = self.unmerged_meta[subproc_index]["flavors"][ps_flavor] - flavors.append(flav["options"][0]) - flavor_remap.append(flav["index"]) - flavor_factors.append(len(flav["options"])) - flavor_mirror.append(flav["mirror"]) - - cross_sections = [] - self.step_matrix_elements = [] - for me_index, matrix_element in enumerate(self.matrix_elements): - if self.surrogates: - # the surrogate in the integrand, the full matrix element in - # the second pass, with the random numbers passed in - mat = ms.MatrixElement( - self.surrogates[me_index].index(), - self.particle_count, - [ms.MatrixElement.momenta_in, ms.MatrixElement.alpha_s_in, - ms.MatrixElement.flavor_in], - [ms.MatrixElement.matrix_element_out], - 1, - False, - self.process.me_frame, - self.process.incoming_count, - ) - self.step_matrix_elements.append(ms.MatrixElement( - matrix_element, - [ms.MatrixElement.momenta_in, ms.MatrixElement.alpha_s_in, - ms.MatrixElement.flavor_in, ms.MatrixElement.random_color_in, - ms.MatrixElement.random_helicity_in, - ms.MatrixElement.random_diagram_in], - [ms.MatrixElement.matrix_element_out, - ms.MatrixElement.color_index_out, - ms.MatrixElement.helicity_index_out, - ms.MatrixElement.diagram_index_out], - False, - self.process.me_frame, - self.process.incoming_count, - )) - elif matrix_element: - mat = ms.MatrixElement( - matrix_element, - ms.Integrand.matrix_element_inputs, - ms.Integrand.matrix_element_outputs, - True, - self.process.me_frame, - self.process.incoming_count, - ) - else: - #TODO: not working in merged mode - mat = ms.MatrixElement( - 0xBADCAFE, - self.particle_count, - ms.Integrand.matrix_element_inputs, - ms.Integrand.matrix_element_outputs, - self.meta["diagram_count"], - True, - self.process.me_frame, - self.process.incoming_count, - ) - pdf_grid = None if self.process.leptonic else self.process.pdf_grid - pdf_arg = None if self.process.leptonic else ms.CachedPdf() - cross_sections.append( - ms.DifferentialCrossSection( - matrix_element=mat, - # 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(), - pid_options=[], - pdf1=pdf_arg, - pdf2=pdf_arg, - input_momentum_fraction=True, - decay=self.process.is_decay, - return_matrix_element=bool(self.surrogates), - ) - ) - partial_weights = self.process.systematics_enabled() - madnis_args = self.process.run_card["madnis"] - custom_bias = self.process.custom_function("bias", self.particle_count, flavors) - integrands = [] - for channel in phasespace.channels: - integrands.append(ms.Integrand( - channel.phasespace_mapping, - cross_sections, - channel.adaptive_mapping, - channel.discrete_sym, - channel.discrete_flavor, - flavors, - pdf_grid, - self.process.running_coupling, - self.scale, - phasespace.prop_chan_weights, - phasespace.subchan_weights, - phasespace.cwnet, - phasespace.first_chan_weight_remap, - phasespace.first_remapped_chan_count, - phasespace.second_chan_weight_remap, - phasespace.second_remapped_chan_count, - madnis_training, - drop_cuts_and_rescale, - partial_weights, - channel.channel_weight_indices, - channel.active_flavors, - flavor_remap, - flavor_factors, - flavor_mirror, - flavor_diff_xs_indices, - flavor_subproc_indices, - flavor_per_subproc_remap, - madnis_args["compressed_channel_weight_count"], - custom_bias, - )) - #print(integrands[1].function()) - #for i in integrands: print(i.function()) - return integrands - - -_ROOTED_OPTIONS = ('lhapdf', 'lhapdf_py3', 'lhapdf_py2', 'heptools_install_dir') - -def load_mg5_options(me_dir=None) -> dict: - """Read the tool paths from the MadGraph configuration, so the launcher - knows where LHAPDF and the optional programs (Pythia8/Delphes/MadSpin/ - reweight/analysis) live. - - Files are read least- to most-specific, each overriding the previous, which - is the same layering CommonRunCmd.set_configuration uses: the card inside - this process directory has the last word, then this installation, then the - per-user file. Relative values are resolved against the root of the file - they came from. - - ``me_dir`` is the process directory whose Cards/me5_configuration.txt has - the last word. It defaults to the working directory, which is right for - bin/generate_events (it chdirs there first) but not for an in-process - caller, which has not moved yet. - """ - - import madgraph - mg5dir = os.path.dirname(os.path.dirname(os.path.abspath(madgraph.__file__))) - me_dir = os.getcwd() if me_dir is None else me_dir - - options = { - 'pythia-pgs_path': None, 'pythia8_path': None, - 'madanalysis5_path': None, 'exrootanalysis_path': None, 'delphes_path': None, - 'rivet_path': None, 'contur_path': None, 'f2py_compiler': None, - 'lhapdf': None, 'lhapdf_py3': None, 'lhapdf_py2': None, 'timeout': 0, - 'heptools_install_dir': None, 'hepmc3_path': None, - } - config_files = [] - base_config = misc.base_config_file() - if base_config: - config_files.append((base_config, mg5dir)) - user_config = misc.user_config_file() - if user_config: - config_files.append((user_config, mg5dir)) - config_files.append((misc.install_config_file(mg5dir), mg5dir)) - config_files.append((os.path.join(me_dir, 'Cards', 'me5_configuration.txt'), - me_dir)) - for cfg, root in config_files: - if not os.path.exists(cfg): - continue - with open(cfg) as fsock: - for line in fsock: - line = line.split('#', 1)[0] - if '=' not in line: - continue - name, value = (x.strip() for x in line.split('=', 1)) - if name not in options or value in ('', 'None'): - continue - if (name.endswith('_path') or name in _ROOTED_OPTIONS) and \ - not os.path.isabs(value) and \ - (os.sep in value or value.startswith('.')): - value = os.path.normpath(os.path.join(root, value)) - options[name] = value - options['mg5_path'] = mg5dir # enables MadSpin/reweight - return options - - -_LHAPDF = None - -def lhapdf_paths(refresh: bool = False) -> misc.LhapdfPaths: - """The LHAPDF installation this run should use, resolved once. - - Uses the same helper as the 'launch' command, so driving a process - directory straight through bin/generate_events behaves identically to - launching it from MadGraph. - """ - - global _LHAPDF - if _LHAPDF is None or refresh: - _LHAPDF = misc.resolve_lhapdf(load_mg5_options()) - # Export what we found: the real LHAPDF library, systematics and the - # MadSpin/reweight subprocesses read these from the environment. - if _LHAPDF.data_paths: - os.environ["LHAPDF_DATA_PATH"] = os.pathsep.join(_LHAPDF.data_paths) - if _LHAPDF.config: - os.environ["MADGRAPH_LHAPDF_CONFIG"] = _LHAPDF.config - return _LHAPDF - - -class MG7Cmd(Cmd): - """The mg7 run interface: the command object that drives a launch. - - It plays two roles, and they are the same object on purpose. - - 1. It is the *mother* of the merged switch/card question (the thing - ``AskRunEditCard`` calls back into for ``keep_cards``/``do_open``/ - ``compute_widths``). - 2. It is the child command interface MG5 registers via - ``define_child_cmd_interface`` when you type ``launch`` on an mg7 output. - - Role 2 is what makes role 1 behave. ``Cmd.ask`` resolves a scripted answer - through ``self.check_answer_in_input_file``, where ``self`` is the mother -- - so once MG5 hands us its ``inputfile`` and registers itself as our - ``mother``, the question reads the very lines that follow ``launch`` in the - user's command file, and any line it does not recognise goes back up to MG5 - through ``store_line`` instead of being silently consumed. This is exactly - how a madevent output already behaves; before this class existed the mg7 - branch shelled out to ``bin/generate_events`` and had to guess which - following lines belonged to the run, which quietly swallowed every ``set``. - """ - - prompt = 'MG7> ' - - def __init__(self, me_dir='.', options=None, *args, **opts): - super().__init__(*args, **opts) - self.me_dir = os.path.abspath(me_dir) if me_dir != '.' else '.' - # The rest of the launcher is written against the process directory as - # the working directory, so do_generate_events chdirs there; but at - # construction time MG5 has not moved, hence the explicit me_dir. - self.options = load_mg5_options( - self.me_dir if self.me_dir != '.' else None) - if options: - # MG5's own resolved tool paths win: they are the ones the user - # configured in the session that is doing the launching. - self.options.update({k: v for k, v in options.items() - if k in self.options and v is not None}) - self.plugin_path = [] - self.proc_characteristics = {'grouped_matrix': False, 'limitations': []} - self._model = None # lazily imported by get_model() - - # ------------------------------------------------------------------ - # the contract the card question expects from its mother - # ------------------------------------------------------------------ - def keep_cards(self, need_card=[], ignore=[]): - from madgraph.interface.common_run_interface import CommonRunCmd - return CommonRunCmd.keep_cards(self, need_card, ignore) - - def do_open(self, line): - from madgraph.interface.common_run_interface import CommonRunCmd - CommonRunCmd.do_open(self, line) - - def check_open(self, args): - from madgraph.interface.common_run_interface import CommonRunCmd - CommonRunCmd.check_open(self, args) - - def get_model(self): - """The model this process was generated with (an ``import_ufo`` model). - - This is the hook the card question goes through for everything that - needs the model rather than the card: 'update dependent' (which resets - the masses/widths the model computes from the free parameters) and the - generic auto-width handler both call ``mother_interface.get_model()``. - Without it the launch question could only warn that it had failed to - update the dependent parameters, and MadSpin/the shower then read - whatever inconsistent values the param_card still held. - - None when the process recorded no model (an output made before - SubProcesses/model.txt existed) -- the caller warns in that case. - """ - if self._model is None: - me_dir = self.me_dir if self.me_dir != '.' else None - # False (not None) so a model that cannot be imported is not - # retried at every question. - self._model = load_process_model(me_dir) or False - return self._model or None - - def do_compute_widths(self, line): - # The interactive card editor delegates 'auto' width computation to - # the mother interface. Reuse the runtime helper (madgraph subprocess - # + the model stored at output time). ``line`` looks like - # " --path= [--nlo]"; we only need the card path. - m = re.search(r'--path=(\S+)', line or "") - path = m.group(1) if m else os.path.join("Cards", "param_card.dat") - compute_auto_widths(path) - # return an empty mapping: the caller iterates out.items() for the - # small-width treatment, which mg7 does not apply. - return {} - - # ------------------------------------------------------------------ - # running - # ------------------------------------------------------------------ - def do_generate_events(self, line): - """generate_events [-f] [-n NAME] [--laststep=parton] - - Run the mg7 generation and the selected post-processing tools.""" - - opts = self._parse_run_options(line) - cwd = os.getcwd() - target = self.me_dir if self.me_dir != '.' else cwd - try: - os.chdir(target) - _setup_logging() - if opts['name']: - self._set_run_name(opts['name']) - switch = {} - if not opts['force']: - switch = ask_edit_cards(mother=self) - switch = self._apply_laststep(switch, opts) - force_lhe_output_if_needed(switch) - _raise_open_file_limit() - run_generation(switch) - finally: - os.chdir(cwd) - - # ``launch`` is the name MG5 users type; keep it working here too. - do_launch = do_generate_events - - def do_npy_to_lhe(self, line): - """npy_to_lhe [RUN_NAME] [-o OUTPUT] [--seed N] [--no-gzip] - - Write the LHE file (events.lhe.gz) of a run that produced npy events, - e.g. to shower it. Without a run name, the most recent such run.""" - import shlex - me_dir = self.me_dir if self.me_dir != '.' else os.getcwd() - try: - npy_to_lhe.main(shlex.split(line), me_dir=me_dir) - except SystemExit as error: - if error.code not in (None, 0): - logger.error(str(error.code)) - - def do_quit(self, line): - """Leave the mg7 run interface.""" - return super().do_quit(line) - - do_exit = do_quit - - # ------------------------------------------------------------------ - # helpers - # ------------------------------------------------------------------ - def _parse_run_options(self, line): - """Understand the subset of MG5's `launch` options that mean something - for an mg7 run, and say so out loud for the ones that do not.""" - out = {'force': False, 'name': '', 'laststep': ''} - args = self.split_arg(line) - index = 0 - while index < len(args): - arg = args[index] - index += 1 - if arg in ('-f', '--force'): - out['force'] = True - elif arg in ('-n', '--name') and index < len(args): - out['name'] = args[index] - index += 1 - elif arg.startswith(('-n=', '--name=')): - out['name'] = arg.split('=', 1)[1] - elif arg in ('-s', '--laststep') and index < len(args): - out['laststep'] = args[index] - index += 1 - elif arg.startswith(('-s=', '--laststep=')): - out['laststep'] = arg.split('=', 1)[1] - elif arg.startswith('-'): - logger.warning( - "'%s' is not supported by the mg7 output and is ignored; " - "the equivalent settings live in Cards/run_card.toml", arg) - elif arg and not out['name']: - out['name'] = arg - return out - - def _set_run_name(self, name): - """Honour `launch -n NAME`: the mg7 run directory is - Events/_NN, so the name is a run_card setting.""" - path = os.path.join("Cards", "run_card.toml") - run_card = RunCardMG7(path, consistency=False) - if run_card["run"]["run_name"] != name: - run_card["run"]["run_name"] = name - run_card.write(path) - logger.info("run name set to '%s'", name) - - def _apply_laststep(self, switch, opts): - """`--laststep=parton` means: generate the events, run nothing on them. - - The mg7 tool selection is the switch dict the question returns, so - stopping at parton level is switching those entries off. - """ - if opts['laststep'] not in ('parton', 'auto', ''): - logger.warning( - "--laststep=%s is not supported by the mg7 output; " - "select the programs to run in the launch question instead", - opts['laststep']) - return switch - if opts['laststep'] != 'parton' or not switch: - return switch - switch = dict(switch) - for key in ('shower', 'detector', 'analysis', 'madspin', 'reweight'): - if not _off(switch.get(key, 'OFF')): - logger.info("--laststep=parton: not running %s", key) - switch[key] = 'OFF' - return switch - - -def build_selector_cmd(mother=None): - """Build the (monkey-patched) mother command + merged switch/card selector - used by the mg7 output. Returns the selector *class* and a mother instance - understood by AskRun/AskforEditCard. - - ``mother`` lets an in-process caller supply its own :class:`MG7Cmd` -- the - one MG5 registered as its child -- so that the question resolves scripted - answers against MG5's own command file (see :class:`MG7Cmd`). When it is - omitted a standalone instance is built, which is what ``bin/generate_events`` - gets. - """ - - from madgraph.interface.madevent_interface import AskRunEditCard - - from madgraph.various import banner as _banner_mod - from madgraph.various import misc as _misc - - class MG7Selector(AskRunEditCard): - """Merged switch/card question for the mg7 output. - - The mg7 run_card is a TOML file, so define_paths/init_run/do_set are - overridden *on this class* (not globally on AskforEditCard) to treat it - as such. Scoping them here is important: MadSpin/reweight spawn their own - legacy madevent runs in the same process, and a global patch would break - their (run_card.dat based) card editing.""" - - # param_card + the TOML run_card are always offered - always_cards = ['param_card.dat', 'run_card.toml'] - optional_cards = [] - - def define_paths(self, **opt): - super().define_paths(**opt) - self.paths["run"] = os.path.join(self.me_dir, "Cards", "run_card.toml") - self.paths["run_card.toml"] = os.path.join(self.me_dir, "Cards", "run_card.toml") - # the TOML run_card uses its own default file (concrete defaults - # written at output time); this powers "set default". - self.paths["run_default"] = os.path.join(self.me_dir, "Cards", "run_card_default.toml") - - def init_run(self, cards): - # Make sure the run_card is loaded as a RunCardMG7 (the generic - # editor does not recognise run_card.toml), else "set " fails. - out = super().init_run(cards) - if not isinstance(getattr(self, "run_card", None), RunCardMG7): - toml_path = self.paths.get("run") or os.path.join( - self.me_dir, "Cards", "run_card.toml") - if os.path.exists(toml_path): - try: - # allow_scan so a run_card holding scan:[...] values loads - with _misc.TMP_variable(_banner_mod.RunCard, "allow_scan", True): - self.run_card = RunCardMG7(toml_path, consistency="warning") - self.run_set = list(self.run_card.keys()) - except Exception as err: - logger.warning("could not load %s: %s", toml_path, err) - if isinstance(getattr(self, "run_card", None), RunCardMG7): - self.run_card.allow_scan = True - return getattr(self, "run_set", out) - - def do_set(self, line, *args, **kwargs): - # madevent-style shortcuts (lhc/lep/fixed_scale/no_parton_cut), cut - # editing, energy units and arithmetic/mass expressions on the TOML - # run_card, intercepted before delegating to the generic editor. - targs = self.split_arg(line) - run_card = getattr(self, "run_card", None) - if isinstance(run_card, RunCardMG7) and targs: - start = 1 if targs[0] == "run_card" else 0 - if len(targs) > start: - name = targs[start] - nlow = name.lower() - rest = " ".join(targs[start + 1:]).split("#")[0].strip() - masses = run_card.get_mass_shortcuts(getattr(self, "param_card", None)) - - if nlow in ("no_parton_cut", "nocut", "no_cut"): - run_card.remove_all_cut() - logger.info("removing all cuts from the run_card.toml") - self.modified_card.add("run") - return - # [histograms] is a list of observables, not a parameter, - # so "set histograms ..." can only mean the whole section. - if nlow in ("histograms", "histogram", "plots"): - value = rest.lower() - if value in ("off", "none", "no", "false", "0"): - run_card.remove_all_histograms() - logger.info( - "removing all histograms from the run_card.toml") - self.modified_card.add("run") - return - if value in ("default", "on"): - if run_card.restore_default_histograms( - self.paths.get("run_default")): - logger.info( - "restored the %d histograms written at " - "output time", - len(run_card["histograms"])) - self.modified_card.add("run") - else: - logger.warning( - "no run_card_default.toml to restore the " - "histograms from") - return - logger.warning( - "'set histograms %s': only OFF (remove them all) " - "and default (put the generated ones back) are " - "understood; edit the [histograms] section of " - "Cards/run_card.toml for anything else", - rest or "(no value)") - return - if nlow in ("lhc", "lep", "ilc", "lcc") and rest: - ecm = run_card.set_collider(nlow, rest, masses) - logger.info("set %s collider: e_cm = %s GeV", nlow, ecm) - self.modified_card.add("run") - return - if nlow == "fixed_scale" and rest: - val = run_card.set_fixed_scale(rest, masses) - logger.info("set fixed scales to %s GeV", val) - self.modified_card.add("run") - return - - if rest and run_card.is_cut_name(name): - cut, bound, val = run_card.set_cut(name, run_card.evaluate(rest, masses)) - logger.info("modify cut %s.%s of the run_card.toml to %s", cut, bound, val) - self.modified_card.add("run") - return - - # 'iseed' is the madevent spelling of the mg7 run.seed, - # but the two disagree on how to ask for a random seed: - # madevent uses 0, mg7 uses -1 (0 being a perfectly good - # fixed seed there). Translate here rather than in - # _LO_SCALAR_MAP, which also drives the LO -> MG7 run_card - # conversion and would silently pin those runs to seed 0. - if rest and nlow == 'iseed': - value = run_card.evaluate(rest, masses) - try: - value = int(value) - except (TypeError, ValueError): - value = None - if value is None: - logger.warning("ignoring 'set iseed %s': not an integer", rest) - return - seed = -1 if value == 0 else value - run_card.set('run.seed', seed, user=True) - logger.info("set iseed (mg7 run.seed) of the run_card.toml to %s%s", - seed, " (random)" if seed == -1 else "") - self.modified_card.add("run") - return - - # legacy madevent run_card names -> mg7 "section.key", so a - # madevent-style launch script ("set nevents 500", "set - # use_syst F", "set bwcutoff 10", ...) edits the mg7 - # run_card.toml verbatim. RunCardMG7._LO_SCALAR_MAP is the - # same rename table used by the LO->MG7 run_card conversion. - if rest and nlow in run_card._LO_SCALAR_MAP \ - and nlow not in [k.lower() for k in run_card.keys()]: - target = run_card._LO_SCALAR_MAP[nlow] - run_card.set(target, rest, user=True) - logger.info("set %s (mg7 %s) of the run_card.toml to %s", - nlow, target, run_card[target]) - self.modified_card.add("run") - return - - if rest and nlow in [k.lower() for k in run_card.keys()]: - current = run_card[nlow] - if isinstance(current, (int, float)) and not isinstance(current, bool): - resolved = run_card.evaluate(rest, masses) - if not isinstance(resolved, str): - prefix = "run_card " if start == 1 else "" - line = "%s%s %s" % (prefix, name, resolved) - return super().do_set(line, *args, **kwargs) - - return MG7Selector, (mother if mother is not None else MG7Cmd()) - - -def ask_edit_cards(mother=None) -> dict: - """Single (MadDM-style) question letting the user both pick which programs - to run after generation and edit the associated cards. Returns the switch - dict describing the selected tools. - - ``mother`` is the :class:`MG7Cmd` asking the question. Passing the instance - MG5 registered as its child is what lets a scripted ``launch`` answer this - question from MG5's own command file. - """ - - selector_class, mother = build_selector_cmd(mother) - # path_msg is what makes Cmd.check_answer_in_input_file accept a bare path as - # an answer (its "elif path:" branch), so that a scripted launch can hand the - # question a card/banner path -- as the question itself advertises -- and have - # it replace the corresponding card. Without it the path is rejected ("This - # answer is not valid for current question") and the default is used instead. - try: - switch, question = mother.ask('', '0', [], path_msg='enter path', - ask_class=selector_class, - mode='auto', line_args=[], force=False, - return_instance=True) - except BaseException: - # interrupted (e.g. ctrl-C) before the prune below could run: prune - # here too, using the instance ask() stashed pre-interrupt, or the - # eagerly-materialised tool cards get misread as "left ON" next time. - instance = getattr(mother, '_last_ask_instance', None) - if instance is not None and hasattr(instance, 'active_cards'): - try: - mother.keep_cards(instance.active_cards()) - except Exception as error: - logger.debug('could not prune tool cards after aborted launch question: %s', error) - raise - switch = dict(switch) - prune_unselected_tool_cards(mother, question, switch) - return switch - - -def prune_unselected_tool_cards(mother, question, switch) -> None: - """Hide the cards of the tools that were NOT selected in the question. - - The question materialises every candidate tool card (copying the - *_default.dat) so that it can offer them all for edition; madevent's - ask_run_configuration prunes them again afterwards (keep_cards on the - selected cards), but the mg7 launcher had no such cleanup. The materialised - cards therefore persisted, and a later re-launch of the same output saw them - all present and defaulted every switch to ON (set_default_ keys off - card presence). Prune here too, using the question's own switch->card map, so - a re-launch reflects the previous selection instead of turning everything on. - """ - try: - keep = list(question.always_cards) - for spec in question.switch_cards: - if spec['on'](switch): - keep.append(spec['card']) - mother.keep_cards(keep) - except Exception as error: - logger.debug('could not prune unselected tool cards: %s', error) - - -def _find_event_file(run_path): - """Locate the LHE event file produced by generate_events (if any).""" - for name in ("events.lhe", "events.lhe.gz"): - path = os.path.join(run_path, name) - if os.path.exists(path): - return path - return None - - -def _find_or_convert_event_file(run_path, log): - """The run's LHE event file, converting its npy events when that is all - the run wrote (see npy_to_lhe).""" - lhe_path = _find_event_file(run_path) - if lhe_path is None and npy_to_lhe.can_convert(run_path): - log.info("No LHE event file in %s: converting the npy events.", run_path) - try: - lhe_path = npy_to_lhe.convert(run_path) - except Exception as error: - _report_failure(log, "npy->LHE conversion", error, run_path) - return lhe_path - - -def _report_failure(log, what, error, directory=None): - """Log a post-processing failure and write the full traceback to a file - (whose path is printed) so the problem can be investigated.""" - import traceback - if not directory or not os.path.isdir(directory): - directory = os.getcwd() - path = os.path.join(directory, "%s_crash.log" % what.replace(" ", "_")) - try: - with open(path, "w") as fsock: - fsock.write(traceback.format_exc()) - except Exception: - path = None - log.warning("%s failed: %s", what, error) - if path: - log.warning("full traceback written to: %s", path) - - -def _off(value) -> bool: - """True when a switch value means the tool was not selected.""" - return value in (None, "OFF", "Not Avail.", "Not Avail. (numpy missing)") - - -def _has_effective_handler(lg) -> bool: - """True when ``lg`` (or an ancestor it propagates to) already has a handler, - i.e. some other party is already displaying these records.""" - while lg: - if lg.handlers: - return True - if not lg.propagate: - return False - lg = lg.parent - return False - - -_TOOL_LOGGING_READY = False - - -def _setup_logging(): - """Make the reused madevent tool drivers' progress visible on screen. - - The drivers already narrate what they do through logger.info / - update_status ("Running Pythia8 [arXiv:...]", "Splitting .lhe event - file...", the live Idle/Running/Completed job counters, "Running - MadSpin", ...), but the standalone mg7 launcher never attaches a handler - to the 'madgraph'/'madevent' loggers, so all of it is swallowed. Attach a - colored INFO console handler to them (idempotent).""" - global _TOOL_LOGGING_READY - if _TOOL_LOGGING_READY: - return - if _has_effective_handler(logging.getLogger("madgraph")): - # Someone already configured logging for us -- in practice MG5, which - # now runs the launch in its own process. Its configuration is - # authoritative (it owns stdout_level and the tutorial logger); adding - # a second handler here would print every line twice. - _TOOL_LOGGING_READY = True - return - try: - import madgraph.interface.coloring_logging # registers ColorFormatter - formatter = logging.ColorFormatter("%(message)s") - except Exception: - formatter = logging.Formatter("%(message)s") - handler = logging.StreamHandler(sys.stdout) - handler.setFormatter(formatter) - handler._mg7_tool_handler = True - for name in ("madgraph", "madevent", "cmdprint", "madgraph7"): - lg = logging.getLogger(name) - if not any(getattr(h, "_mg7_tool_handler", False) for h in lg.handlers): - lg.addHandler(handler) - lg.setLevel(logging.INFO) - lg.propagate = False - _TOOL_LOGGING_READY = True - - -def selected_tools(switch) -> list: - """The post-processing drivers run_selected_tools runs for this switch.""" - # MadAnalysis5 hadron level analyses the shower/detector output, so it only - # makes sense when a shower ran (mirrors madevent's card gating: - # analysis == 'MadAnalysis5' and shower != 'OFF'). - ma5 = switch.get("analysis") == "MadAnalysis5" - showered = not _off(switch.get("shower")) - return [t for t, on in ( - ("reweighting", not _off(switch.get("reweight"))), - ("MadSpin", not _off(switch.get("madspin"))), - ("MadAnalysis5 (parton level)", ma5), - ("Pythia8 shower", switch.get("shower") == "Pythia8"), - ("Delphes", switch.get("detector") == "Delphes"), - ("MadAnalysis5 (hadron level)", ma5 and showered), - ("Rivet", switch.get("analysis") == "Rivet"), - ) if on] - - -def run_selected_tools(switch, process) -> None: - """Run the optional post-processing programs selected in the merged - question on the generated events. - - All the tools are driven through :class:`MG7RunCmd`, a thin adapter around - the standard madevent run interface (``MadEventCmd``). This means the mg7 - output reuses the *exact* madevent tool drivers -- in particular Pythia8 is - showered in parallel over the cluster/multicore backend, identically to a - madevent run -- instead of re-implementing each tool here. The command - sequence mirrors ``MadEventCmd.do_launch`` (reweight, MadSpin, MA5 parton, - shower, Delphes, MA5 hadron, Rivet); each tool checks its own card, so a - tool is invoked only when its switch is on. - """ - log = logging.getLogger("madevent") - - tools = selected_tools(switch) - if not tools: - # A switch that is not off does not necessarily select a driver here: - # "Not Avail." is not off, analysis = ExRoot has no mg7 driver, and a - # shower switch set to something other than Pythia8 selects nothing - # either. Without this a plain generate/output/launch announced a - # post-processing step with nothing after the colon, and paid for - # building the run interface (or an npy->LHE conversion) to do nothing. - return - - lhe_path = _find_or_convert_event_file(process.run_path, log) - if lhe_path is None: - log.warning("No LHE event file in %s; cannot run %s.", - process.run_path, ", ".join(tools)) - return - - run_name = os.path.basename(os.path.dirname(os.path.abspath(lhe_path))) - run_dir = os.path.dirname(os.path.abspath(lhe_path)) - - ma5 = switch.get("analysis") == "MadAnalysis5" - showered = not _off(switch.get("shower")) - - log.info("") - log.info("Post-processing the generated events with: %s", ", ".join(tools)) - - try: - from madgraph.iolibs.template_files.mg7.run_interface import MG7RunCmd - cmd = MG7RunCmd(os.getcwd(), load_mg5_options(), run_name, lhe_path) - cmd.set_run_name(run_name, None, "parton") - except Exception as error: - _report_failure(log, "post-processing setup", error, run_dir) - return - - def run(tool, command): - bar = "=" * 60 - log.info("") - log.info(bar) - log.info(" post-processing step: %s", tool) - log.info(bar) - start = time.time() - try: - cmd.exec_cmd(command, postcmd=False, printcmd=False) - except Exception as error: - _report_failure(log, tool, error, run_dir) - else: - log.info(" -> %s done (%.1fs)", tool, time.time() - start) - - # Same order as MadEventCmd.do_launch. The commands take no run name and so - # act on cmd.run_name: reweight runs on the generated events, then - # decay_events (MadSpin) repoints cmd.run_name to the "_decayed_i" run, - # so the shower and the analyses that follow act on the decayed events -- - # exactly as a madevent run does. Each driver still guards on its own card. - if not _off(switch.get("reweight")): - run("reweight", "reweight -from_cards") - if not _off(switch.get("madspin")): - run("MadSpin", "decay_events -from_cards") - if ma5: - run("MadAnalysis5 (parton)", "madanalysis5_parton --no_default") - if switch.get("shower") == "Pythia8": - run("Pythia8 shower", "shower --no_default") - if switch.get("detector") == "Delphes": - run("Delphes", "delphes --no_default") - # hadron-level MA5 needs the shower/detector output: only run it if a shower - # was requested (otherwise there is no hadron-level event file to analyse). - if ma5 and showered: - run("MadAnalysis5 (hadron)", "madanalysis5_hadron --no_default") - if switch.get("analysis") == "Rivet": - run("Rivet", "rivet --no_default") - - # Finalize the run exactly as MadEventCmd.do_launch does after the - # shower/detector/analysis tools: store_result() processes the deferred - # "to_store" actions -- in particular it gzips the Pythia8 HepMC to - # _pythia8_events.hepmc.gz (the pythia8_card default is HEPMCoutput:file - # = hepmc.gz). The deferred Rivet job reads exactly that .gz path, so this - # must run before the Rivet/Contur postprocessor below. - try: - cmd.store_result() - except Exception as error: - _report_failure(log, "store_result", error, run_dir) - - if switch.get("analysis") == "Rivet": - # do_rivet only *prepared* the run (wrote Events//run_rivet.sh) and - # deferred execution to the postprocessor -- the rivet_card default has - # run_rivet_later = True, so with --no_default it logged "Skipping Rivet - # for now, passing it to postprocessor" and appended the run to - # cmd.postprocessing_dirs. MadEventCmd.do_launch runs that postprocessor - # at the end of a run; do the same here so the .yoda (and, when - # run_contur = True, the Contur limits) are actually produced. - bar = "=" * 60 - log.info("") - log.info(bar) - log.info(" post-processing step: Rivet/Contur postprocessing") - log.info(bar) - start = time.time() - try: - cmd.postprocessing() - except Exception as error: - _report_failure(log, "Rivet-Contur postprocessing", error, run_dir) - else: - log.info(" -> Rivet/Contur postprocessing done (%.1fs)", - time.time() - start) - - -def _add_time_of_flight(lhe_path, threshold, param_card_path, log): - """Add invariant-lifetime (vtim) information to the LHE events, drawing each - unstable particle's decay length from its width (in mm). Mirrors - common_run_interface.CommonRunCmd.do_add_time_of_flight but reads the widths - from the local param_card instead of the LHE banner.""" - import random - try: - import madgraph.various.lhe_parser as lhe_parser - except ImportError: - import internal.lhe_parser as lhe_parser - from models import check_param_card as param_card_mod - from madgraph.various import misc as _misc - import madgraph.iolibs.files as files - - need_zip = lhe_path.endswith('.gz') - if need_zip: - _misc.gunzip(lhe_path) - lhe_path = lhe_path[:-3] - - param_card = param_card_mod.ParamCard(param_card_path) - cst = 6.58211915e-25 # hbar in GeV s - c = 299792458000 # speed of light in mm/s - log.info('Adding time of flight information on %s', lhe_path) - lhe = lhe_parser.EventFile(lhe_path) - out = open('%s_2vertex.lhe' % lhe_path, 'w') - out.write(lhe.banner) - for event in lhe: - for particle in event: - # default=0 -> particles without a decay entry are treated as stable - width = param_card['decay'].get((abs(particle.pid),), 0.).value - if width: - vtim = c * random.expovariate(width / cst) - if vtim > threshold: - particle.vtim = vtim - out.write(str(event)) - out.write('\n') - out.close() - lhe.close() - files.mv('%s_2vertex.lhe' % lhe_path, lhe_path) - if need_zip: - _misc.gzip(lhe_path) - - -def _lhapdf_config_path(): - """Path to lhapdf-config so systematics can import the python lhapdf - module (required to compute PDF/scale variations).""" - return lhapdf_paths().config - - -def _run_systematics(lhe_path, cfg, log): - """Run systematics.py (scale/PDF variations) on the LHE file, reusing the - exact worker madevent uses (systematics.call_systematics).""" - try: - import madgraph.various.systematics as systematics - except ImportError: - import internal.systematics as systematics - - def fmt(vals): - return ','.join(str(v) for v in vals) - - opts = [] - if cfg.get('systematics_mur'): - opts.append('--mur=%s' % fmt(cfg['systematics_mur'])) - if cfg.get('systematics_muf'): - opts.append('--muf=%s' % fmt(cfg['systematics_muf'])) - if cfg.get('systematics_pdf'): - opts.append('--pdf=%s' % fmt(cfg['systematics_pdf'])) - extra = cfg.get('systematics_str_options', '') - if extra: - opts.extend(extra.split()) - # The mg7 LHE has no block, so systematics cannot read the per-event - # LO reweighting info. For a process with a single alpha_s power (stored in - # SubProcesses/proc_characteristics at output time) it can be reconstructed - # from the events: pass that power as --lo_nqcd. -1 means the QCD power is - # not uniform, so the reconstruction is not applicable. - if not any(o.startswith('--lo_nqcd') for o in opts): - try: - from madgraph.various import banner as _banner_mod - except ImportError: - import internal.banner as _banner_mod - pc_path = os.path.join('SubProcesses', 'proc_characteristics') - if os.path.exists(pc_path): - try: - nqcd = int(_banner_mod.ProcCharacteristic(pc_path)['single_qcd_order']) - except Exception: - nqcd = -1 - if nqcd >= 0: - opts.append('--lo_nqcd=%d' % nqcd) - # tell systematics where to find lhapdf (so it can link the python module) - if not any(o.startswith('--lhapdf_config') for o in opts): - lhapdf_config = _lhapdf_config_path() - if lhapdf_config: - opts.append('--lhapdf_config=%s' % lhapdf_config) - - log.info('Running systematics on %s %s', lhe_path, ' '.join(opts)) - systematics.call_systematics([lhe_path, lhe_path] + opts, - log=lambda x: log.info(str(x))) - - -def run_lhe_postprocessing(process) -> None: - """Run the LHE-level post-processings configured in the run_card - [postprocessing] section on the generated event file (displaced-vertex - time-of-flight and systematics). Only applies when an LHE file exists.""" - lhe_path = _find_event_file(process.run_path) - if lhe_path is None: - return - # process.run_card is a RunCardMG7; the section view exposes .get(key, def) - try: - cfg = process.run_card["postprocessing"] - except Exception: - return - log = logging.getLogger('madevent') - - 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'): - native = getattr(process, 'systematics', None) - if native is not None and native.weight_count > 0: - log.warning("[postprocessing] systematics is ignored: the scale/PDF " - "weights were already computed by madspace ([systematics] " - "section); set systematics.enable = false to use the legacy " - "systematics.py path instead.") - else: - try: - _run_systematics(lhe_path, cfg, log) - except Exception as error: - _report_failure(log, "systematics computation", error, - os.path.dirname(lhe_path)) - - tof = cfg.get('time_of_flight', -1.0) - try: - if tof is not None and float(tof) >= 0: - _add_time_of_flight(lhe_path, float(tof), process.param_card_path, log) - except Exception as error: - _report_failure(log, "add_time_of_flight", error, - os.path.dirname(lhe_path)) - - -_model_hash_warned = set() -_model_cache = {} - - -def read_stored_model(me_dir=None) -> "str | None": - """The model this process was generated with, as recorded at output time in - ``SubProcesses/model.txt``: a UFO path carrying its restriction (e.g. - ``.../models/sm-no_b_mass``) or a model name -- in either case something - ``import model`` understands. Returns None when the process was written - without that information. - - Warns (once per directory) when the model on disk no longer matches the - hash stored on the second line, i.e. it changed since the output was made. - """ - base = me_dir or os.getcwd() - model_file = os.path.join(base, "SubProcesses", "model.txt") - if not os.path.exists(model_file): - return None - with open(model_file) as f: - lines = f.read().splitlines() - model = lines[0].strip() if lines else "" - stored_hash = lines[1].strip() if len(lines) > 1 else "" - if not model: - logger.warning("SubProcesses/model.txt is empty; the model of this " - "process is unknown.") - return None - - # verify the model on disk still matches the one used at output time. The - # reference may carry a restriction suffix ('sm-no_b_mass'); the hash was - # taken on the UFO directory itself ('sm'). - model_dir = model if os.path.isdir(model) else model.rsplit('-', 1)[0] - if stored_hash and os.path.isdir(model_dir) \ - and model_file not in _model_hash_warned: - current_hash = misc.hash_model_files(model_dir) - if current_hash and current_hash != stored_hash: - _model_hash_warned.add(model_file) - logger.warning( - "The model at %s has changed since this process was generated " - "(hash mismatch); anything computed from it now (the 'auto' " - "widths, the dependent masses/widths of the param_card) may be " - "inconsistent with the matrix element.", model_dir) - return model - - -def _proc_characteristic(me_dir=None) -> dict: - """SubProcesses/proc_characteristics as a plain dict (empty when the file is - missing or unreadable). Plain, because ProcCharacteristic is a ConfigFile, - whose get() is the parameter accessor and takes no default.""" - from madgraph.various import banner as _banner_mod - path = os.path.join(me_dir or os.getcwd(), 'SubProcesses', - 'proc_characteristics') - if os.path.exists(path): - try: - return dict(_banner_mod.ProcCharacteristic(path)) - except Exception as error: - logger.debug("could not read %s: %s", path, error) - return {} - - -def load_process_model(me_dir=None): - """Import the UFO model this process was generated with, as - :func:`models.import_ufo.import_model` builds it -- the object every - model-level consumer of the param_card expects (see - :meth:`MG7Cmd.get_model`). None when the process recorded no model, or when - that model no longer imports. - - Cached per (reference, complex-mass-scheme): importing a model is slow and - the card question asks for it again at every 'update dependent'.""" - ref = read_stored_model(me_dir) - if not ref: - return None - cms = _proc_characteristic(me_dir).get('complex_mass_scheme', False) - key = (ref, bool(cms)) - if key not in _model_cache: - import models.import_ufo as import_ufo - try: - with misc.MuteLogger(['madgraph.model'], [50]): - _model_cache[key] = import_ufo.import_model( - ref, complex_mass_scheme=bool(cms)) - except Exception as error: - # 'update dependent' calls this under a SIGALRM whose handler - # raises a TimeOutError of its own (a class defined inside the - # caller, so it can only be recognised by name). Let it through: - # the caller then says the model took too long to load and how to - # force it, instead of reporting a model that does not import. - if type(error).__name__ == 'TimeOutError': - raise - logger.debug("could not import the model %s: %s", ref, error) - _model_cache[key] = None - return _model_cache[key] - - -def compute_auto_widths(param_card_path=os.path.join("Cards", "param_card.dat")) -> None: - """Fill any width set to ``auto`` in the param_card, using madgraph and the - model stored at output time (``SubProcesses/model.txt``), and write the - result back into the card. A no-op when the card has no ``auto`` width. - - Called before every generation, so a single run and each scan point (whose - param_card is rewritten just before ``run_single``) both get model-computed - widths.""" - try: - with open(param_card_path) as f: - text = f.read() - except OSError: - return - # matches "DECAY auto" (optionally "auto@NLO"), as in - # common_run_interface.static_check_param_card - pdgs = re.findall(r"(?im)^\s*decay\s+([+-]?\d+)\s+auto", text) - if not pdgs: - return - pdgs = list(dict.fromkeys(pdgs)) # de-duplicate, keep order - - model = read_stored_model() - if not model: - logger.warning( - "The param_card requests 'auto' width(s) for %s but the model of " - "this process is unknown; leaving them as-is.", " ".join(pdgs)) - return - - mg5 = str(_MG_ROOT / "bin" / "madgraph") - if not os.path.exists(mg5): - logger.warning("Cannot find madgraph at %s; 'auto' widths not computed.", mg5) - return - - import tempfile - cmds = "import model %s\ncompute_widths %s --path=%s\n" % ( - model, " ".join(pdgs), os.path.abspath(param_card_path)) - with tempfile.NamedTemporaryFile("w", suffix=".mg5", delete=False) as fh: - fh.write(cmds) - cmdfile = fh.name - try: - logger.info("Computing 'auto' width(s) for %s ...", " ".join(pdgs)) - proc = subprocess.run([mg5, cmdfile]) - if proc.returncode != 0: - logger.warning( - "compute_widths returned a non-zero exit code; the param_card " - "may still contain 'auto' entries.") - finally: - try: - os.remove(cmdfile) - except OSError: - pass - - -def _release_channel_generators(process) -> None: - """Drop the per-channel event/weight generators so their intermediate - .npy files are deleted right away. madspace only frees them when the - owning Python objects are collected, and process.event_generator / - process.phasespaces sit in a reference cycle that plain refcounting - never breaks -- only an explicit gc.collect() does, promptly, instead of - leaving it to whenever the cyclic GC next runs on its own.""" - process.event_generator = None - process.phasespaces = None - gc.collect() - - -def run_single(switch=None) -> "MadgraphProcess": - """Run a single generation and return the process (for its result).""" - compute_auto_widths() - process = MadgraphProcess() - process.survey() - process.train_madnis() - process.generate_events() - if process.gridpack_run_mode != "regular": - return process - # run_card-driven LHE post-processing (displaced vertex + systematics) - run_lhe_postprocessing(process) - # run the post-processing tools (Pythia8/Delphes/MadSpin/reweight/analysis) - # selected in the merged question above on the generated events. - if switch: - run_selected_tools(switch, process) - return process - - -def detect_run_scan(run_card_path): - """Return a banner.RunCardIterator if the run_card contains scan:[...] - values, else None.""" - from madgraph.various import banner as banner_mod - from madgraph.various import misc as _misc - with _misc.TMP_variable(banner_mod.RunCard, 'allow_scan', True): - rc = banner_mod.RunCard(run_card_path, consistency=False) - if getattr(rc, 'scan_set', None): - return banner_mod.RunCardIterator(run_card_path) - return None - - -def detect_param_scan(param_card_path): - """Return a ParamCardIterator if the param_card contains scan:[...] values, - else None.""" - if not os.path.exists(param_card_path): - return None - from models import check_param_card as param_card_mod - it = param_card_mod.ParamCardIterator(param_card_path) - for block in it.order: - for param in block: - if isinstance(param.value, str) and param.value.strip().lower().startswith('scan'): - return it - return None - - -def run_scan(iterator, card_path, switch=None) -> None: - """Iterate over all scan points, running a full generation for each and - accumulating the results, then write the scan summary. Works for both the - run_card (RunCardIterator) and the param_card (ParamCardIterator); their - interface (__iter__/write/store_entry/get_next_name/write_summary) is the - same. The scan card is restored afterwards.""" - import tomllib - with open(os.path.join("Cards", "run_card.toml"), "rb") as f: - run_name = tomllib.load(f).get("run", {}).get("run_name", "run") - - from models import check_param_card as param_card_mod - is_param_scan = isinstance(iterator, param_card_mod.ParamCardIterator) - - backup = card_path + ".scan_bak" - shutil.copy(card_path, backup) - try: - for i, point in enumerate(iterator): - point.write(card_path) - logger.info("=== scan point %d ===", i + 1) - process = run_single(switch) - # use the run directory the process actually created, so the - # per-point params.dat written by write_summary has a home - name = os.path.basename(process.run_path) - if is_param_scan: - # pass the (possibly auto-width-updated) card so the summary - # records the model-computed width for each 'auto' entry - iterator.store_entry(name, process.get_result(), - param_card_path=card_path) - else: - iterator.store_entry(name, process.get_result()) - _release_channel_generators(process) - os.makedirs("Events", exist_ok=True) - summary = os.path.join("Events", "scan_%s.txt" % run_name) - iterator.write_summary(summary) - logger.info("scan results written to %s", summary) - finally: - shutil.move(backup, card_path) - - -def run_generation(switch=None) -> None: - """Run the generation, expanding a scan over the run_card or the param_card - when one is present (scanning both simultaneously is not allowed).""" - run_card_path = os.path.join("Cards", "run_card.toml") - param_card_path = os.path.join("Cards", "param_card.dat") - run_iter = detect_run_scan(run_card_path) - param_iter = detect_param_scan(param_card_path) - if run_iter and param_iter: - raise RuntimeError( - "Scanning simultaneously over the run_card and the param_card is " - "not allowed. Please keep the scan:[...] entries in only one card.") - if run_iter: - run_scan(run_iter, run_card_path, switch) - elif param_iter: - run_scan(param_iter, param_card_path, switch) - else: - _release_channel_generators(run_single(switch)) - - -def _lhe_needed_by(switch, card) -> "str | None": - """Why the run needs an LHE event file, or None when it does not: every - post-processing tool run_selected_tools drives and the run_card - [postprocessing] steps read the LHE file. ``card`` is the raw TOML content - of the run_card.""" - tools = selected_tools(switch) if switch else [] - if tools: - return tools[0] - cfg = card.get("postprocessing", {}) - try: - if float(cfg.get("time_of_flight", -1.0)) >= 0: - return "time_of_flight" - except (TypeError, ValueError): - pass - # the legacy systematics.py path only runs when madspace did not compute - # the weights itself (see run_lhe_postprocessing) - if cfg.get("systematics") \ - and not card.get("systematics", {}).get("enable", True): - return "systematics" - return None - - -def force_lhe_output_if_needed(switch) -> None: - """The events are written as npy by default; switch the run_card to the - LHE format when a selected post-processing needs an LHE file. The card is - rewritten rather than refused, so that asking for a shower in the launch - question is enough -- the user does not have to know that it constrains - output_format. - - The card is read as raw TOML and only its output_format line is edited: a - run_card scan ("scan:[...]" values) is not a valid RunCardMG7 and must - survive untouched for run_scan.""" - import tomllib - path = os.path.join("Cards", "run_card.toml") - if not os.path.exists(path): - return - with open(path, "rb") as f: - card = tomllib.load(f) - output_format = card.get("run", {}).get( - "output_format", RunCardMG7()["run"]["output_format"]) - if output_format == "lhe": - return - reason = _lhe_needed_by(switch, card) - if reason is None: - return - with open(path) as f: - text = f.read() - text, count = re.subn(r'(?m)^(\s*output_format\s*=\s*)("[^"]*"|\'[^\']*\'|\S+)', - r'\g<1>"lhe"', text, count=1) - if not count: - text, count = re.subn(r'(?m)^\[run\][^\n]*\n', '\\g<0>output_format = "lhe"\n', - text, count=1) - if not count: - text += '\n[run]\noutput_format = "lhe"\n' - with open(path, "w") as f: - f.write(text) - logging.getLogger("madevent").info( - "output_format set to 'lhe' (required by %s).", reason) - - -def _raise_open_file_limit() -> None: - """Remove the soft limit on the number of open files: it can be quite low - on some systems and the event-combination step opens one file per channel.""" - soft_lim, hard_lim = resource.getrlimit(resource.RLIMIT_NOFILE) - try: - resource.setrlimit(resource.RLIMIT_NOFILE, (hard_lim, hard_lim)) - except (ValueError, OSError) as error: - logger.debug("could not raise the open-file limit: %s", error) - - -def main() -> None: - """``bin/generate_events`` entry point. - - It runs the same :class:`MG7Cmd` command MG5's in-process ``launch`` uses, - so the two entry points cannot drift apart; the only difference is that - nothing here supplies a ``mother``/``inputfile``, so the question falls back - to reading this process's own stdin. - """ - parser = argparse.ArgumentParser() - parser.add_argument("-f", action="store_false", dest="ask_edit_cards") - args = parser.parse_args() - - MG7Cmd().run_cmd("generate_events" if args.ask_edit_cards else "generate_events -f") diff --git a/madgraph/iolibs/template_files/mg7/madgraph_process.py b/madgraph/iolibs/template_files/mg7/madgraph_process.py new file mode 100644 index 0000000000..0f910aba49 --- /dev/null +++ b/madgraph/iolibs/template_files/mg7/madgraph_process.py @@ -0,0 +1,1995 @@ +"""The run of one process directory: cards, integration, events and output.""" + +import glob +import json +import os +import shutil +import subprocess +import time + +from madgraph.iolibs.template_files.mg7.common import backend_of, format_time, fptype_of, get_start_time, interference_helicities_summed, lhe_helicities, lhe_weight_info, logger, me_parameters, ms, resolve_auto_backend, resolve_seed, resolve_verbosity +from madgraph.various import misc +from madgraph.various.banner import RunCardMG7 +from models.check_param_card import ParamCard +from madgraph.iolibs.template_files.mg7 import hwu_output +from madgraph.iolibs.template_files.mg7 import npy_to_lhe +from madgraph.iolibs.template_files.mg7 import options +from madgraph.iolibs.template_files.mg7 import plots +from madgraph.iolibs.template_files.mg7 import systematics_summary +from madgraph.iolibs.template_files.mg7.bootstrap import INSTALL_DIR as _INSTALL_DIR +from madgraph.iolibs.template_files.mg7.bootstrap import MADSPACE_DIR as _MADSPACE_DIR +from madgraph.iolibs.template_files.mg7.custom_function_loader import load_custom_function +from madgraph.iolibs.template_files.mg7.observables import CutItem, HistItem, build_observable, clean_pids, observable_pids +from madgraph.iolibs.template_files.mg7.options import lhapdf_paths +from madgraph.iolibs.template_files.mg7.subprocess_amplicol import AmplicolSubprocess, amplicol_subproc_args, is_amplicol +from madgraph.iolibs.template_files.mg7.subprocess_madmatrix import MadgraphSubprocess, PhaseSpace, build_multi_channel_data + +class MadgraphProcess: + def __init__(self): + self.madnis_trained = False + self.lhe_completer = None + self.load_cards() + self.init_backend() + self.init_event_dir() + self.init_context() + self.init_custom_functions() + self.init_surrogate() + self.init_cuts() + self.init_histograms() + self.init_generator_config() + self.init_beam() + self.init_subprocesses() + + def load_cards(self) -> None: + self.run_card = RunCardMG7(os.path.join("Cards", "run_card.toml")) + # Resolved once so every generator built during this run shares the same + # seed; the concrete value (even if randomly drawn) is recorded in each + # generator's info.json. + self.run_seed = resolve_seed(self.run_card["run"]["seed"]) + self.param_card_path = os.path.join("Cards", "param_card.dat") + self.param_card = ParamCard(self.param_card_path) + with open(os.path.join("SubProcesses", "subprocesses.json")) as f: + self.subprocess_data = json.load(f) + if self.run_card["phasespace"]["merge_subprocesses"]: + with open(os.path.join("SubProcesses", "merged_subprocesses.json")) as f: + self.merged_subprocess_data = json.load(f) + else: + self.merged_subprocess_data = None + + self.init_decay_mode() + self.init_me_frame() + if interference_helicities_summed(self.run_card): + count = sum(1 for meta in self.subprocess_data if meta.get("interference")) + if count: + logger.info( + "interference_helicity = summed: the weight of the %d interference " + "subprocess(es) is the helicity sum, and their helicities are " + "written as 9 in the LHE", count) + + def init_me_frame(self) -> None: + """Resolve the run card's me_frame into the list of external particles + whose momentum sum defines the frame the matrix element is evaluated in. + + Empty (the default) means no boost: the matrix element sees the momenta + in the frame they are generated in. Unlike madevent, which hands its + matrix element momenta that are already in the partonic centre of mass, + madspace generates them in the lab frame, so asking for madevent's own + frame ([1, 2]) does mean an actual boost here. It changes nothing for a + Lorentz invariant matrix element (only the rounding), but it is what + defines the polarisation axes of a polarised one. + + The particle numbers are checked here rather than only in madspace, + which would not see them until the integrands are built -- long after + the matrix-element libraries have been compiled. Every subprocess is + checked against its own particles: an output can mix multiplicities + (p p > z j and p p > z j j, say, for MLM merging), and madspace builds + one matrix element per subprocess, each refusing a particle it does not + have. That is stricter than madevent, which ignores the particles a + process does not have and so changes the frame from one multiplicity + to the next without a word. + """ + me_frame = list(self.run_card["run"]["me_frame"]) + pdgs = [clean_pids(meta["incoming"]) + clean_pids(meta["outgoing"]) + for meta in self.subprocess_data] + counts = sorted({len(pdg_list) for pdg_list in pdgs}) + for meta, pdg_list in zip(self.subprocess_data, pdgs): + for index in me_frame: + if 1 <= index <= len(pdg_list): + continue + process = "%s > %s" % ( + " ".join(str(pid) for pid in meta["incoming"]), + " ".join(str(pid) for pid in meta["outgoing"]), + ) + mixed = ( + f"; this output mixes subprocesses with " + f"{', '.join(str(count) for count in counts)} external " + "particles, and me_frame has to name particles every one " + "of them has" + if len(counts) > 1 else "" + ) + raise ValueError( + f"me_frame particle {index} out of range: subprocess " + f"{process} has {len(pdg_list)} external particles{mixed}" + ) + if len(me_frame) == 1: + # A massless particle has no rest frame: its momentum squares to + # zero, so the boost divides by it and every momentum comes back as + # a NaN, silently. Name the mistake instead. + massless = [pdg_list[me_frame[0] - 1] for pdg_list in pdgs + if self.get_mass(pdg_list[me_frame[0] - 1]) == 0.] + if massless: + raise ValueError( + f"me_frame = {me_frame} asks for the rest frame of particle " + f"{me_frame[0]}, which is massless (pdg " + f"{sorted(set(massless))}) and has none. Name a massive " + "particle, or several particles whose sum is massive." + ) + self.me_frame = me_frame + self.incoming_count = 1 if self.is_decay else 2 + + 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 and self.merged_subprocess_data is not None: + # merged_subprocesses.json indexes the *unfiltered* subprocess + # list, so dropping entries here would silently shift every index + # it holds. Refuse rather than mis-map. + raise ValueError( + "merge_subprocesses is not supported for a decay directory " + "with kinematically closed channels" + ) + 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"]) + + def init_event_dir(self) -> None: + output_dir = self.run_card["gridpack"]["output_dir"] + temp_dir = self.run_card["gridpack"]["temp_output_dir"] + if output_dir: + self.run_path = os.path.abspath(os.path.expanduser(output_dir)) + os.makedirs(self.run_path, exist_ok=True) + else: + run_name = self.run_card["run"]["run_name"] + os.makedirs("Events", exist_ok=True) + run_dir_prefix = os.path.join("Events", f"{run_name}_") + existing_run_dirs = glob.glob(f"{run_dir_prefix}*") + run_index = 1 + for run_dir in existing_run_dirs: + run_index_str = run_dir[len(run_dir_prefix):] + if run_index_str.isnumeric(): + run_index = max(run_index, int(run_index_str) + 1) + while True: + try: + self.run_path = f"{run_dir_prefix}{run_index:02d}" + os.mkdir(self.run_path) + break + except FileExistsError: + run_index += 1 + # temporary npy files of the channel generators + if temp_dir: + self.temp_path = os.path.abspath(os.path.expanduser(temp_dir)) + os.makedirs(self.temp_path, exist_ok=True) + else: + self.temp_path = self.run_path + # Absolute on purpose. StatusFile is a C++ object that finishes its + # write from its destructor (rename info.json.tmp -> info.json), and + # that destructor runs whenever Python gets round to collecting the + # process object. When the launcher runs in MG5's process, an + # interrupted run is collected *after* the launch command has restored + # the working directory, and a relative path would then resolve + # somewhere else -- the rename throws a C++ filesystem_error out of a + # destructor, which is an immediate abort. + self.status_file = ms.StatusFile( + os.path.abspath(os.path.join(self.run_path, "info.json"))) + + def init_context(self) -> None: + device_names = self.run_card["run"]["device"] + self.contexts = [] + # AmpliCol libraries, shared by all channels of their process + self.amplicol_apis = {} + self.device_types = [] + self.devices = [] + self.pool_sizes = [] + for device_name in device_names: + if ":" in device_name: + device_type, device_index_str = device_name.split(":") + device_index = int(device_index_str) + else: + device_type = device_name + device_index = 0 + self.device_types.append(device_type) + if device_type == "cuda": + device = ms.cuda_device(device_index) + pool_size = self.run_card["run"]["gpu_thread_pool_size"] + elif device_type == "hip": + device = ms.hip_device(device_index) + pool_size = self.run_card["run"]["gpu_thread_pool_size"] + else: + device = ms.cpu_device() + pool_size = self.run_card["run"]["cpu_thread_pool_size"] + self.devices.append(device) + self.pool_sizes.append(pool_size) + self.contexts.append(ms.Context(device=device, thread_count=pool_size)) + + CUSTOM_FUNCTION_KINDS = ("cut", "scale", "pdf", "bias") + + def init_custom_functions(self) -> None: + """Load the functions of [custom_functions], building the libraries + first. They are loaded before the matrix elements, so that they have the + same index in every context and in a gridpack made from this run.""" + self.custom_functions = {} + self.custom_function_paths = [] + for kind in self.CUSTOM_FUNCTION_KINDS: + path = self.run_card["custom_functions"][kind] + if not path: + continue + build_dir = f"custom_{kind}" + if path.endswith(".so") and os.path.isfile(os.path.join(build_dir, "Makefile")): + self.build_custom_function(build_dir) + if not os.path.isfile(path): + raise RuntimeError( + f"[custom_functions] {kind} = {path!r}: file not found. " + "Create a .so library or a .py file with bin/init_custom_func." + ) + apis = load_custom_function(ms, self.contexts, self.device_types, kind, path) + if len({api.index() for api in apis}) != 1: + raise RuntimeError(f"custom {kind} function loaded at different indices") + self.custom_functions[kind] = apis[0] + self.custom_function_paths.append([kind, path]) + logger.info("Using the custom %s function %s", kind, path) + + def init_surrogate(self) -> None: + """Check the [surrogate] settings. The surrogates are loaded per + subprocess, by MadgraphSubprocess.""" + settings = self.run_card["surrogate"] + self.surrogate = None + if not settings["function"]: + return + path = settings["function"] + if path.endswith(".py") and not os.path.isfile(path): + raise RuntimeError(f"[surrogate] function = {path!r}: file not found") + if self.run_card["phasespace"]["sde_strategy"] != "denominators": + raise RuntimeError( + "[surrogate] needs [phasespace] sde_strategy = \"denominators\", " + "as the surrogate does not return channel weights" + ) + if self.run_card["run"]["dummy_matrix_element"]: + raise RuntimeError("[surrogate] cannot be used with dummy_matrix_element") + if self.run_card["gridpack"]["save_gridpack"]: + raise RuntimeError("[surrogate] does not support gridpacks yet") + self.surrogate = dict(settings) + if not self.surrogate["param_card"]: + self.surrogate["param_card"] = self.param_card_path + logger.info( + "Using the surrogate %s, %s with the full matrix element", + path, + "unweighting" if settings["mode"] == "unweight" else "reweighting", + ) + + def weighting_groups(self, subprocess_channels: list[list[int]]) -> list: + """One ms.WeightingGroup per subprocess with weighting steps, over its + channels, given by their indices in the event generator.""" + groups = [] + for subproc, channel_indices in zip(self.subprocesses, subprocess_channels): + steps = subproc.weighting_steps() + if channel_indices and steps: + groups.append(ms.WeightingGroup(channel_indices, steps)) + return groups + + @staticmethod + def build_custom_function(build_dir: str) -> None: + """Run make in a custom_ directory, logging to compile.log.""" + log_path = os.path.join(build_dir, "compile.log") + with open(log_path, "w") as log: + returncode = subprocess.call( + ["make"], cwd=build_dir, stdout=log, stderr=subprocess.STDOUT + ) + if returncode != 0: + raise RuntimeError( + f"Compilation of {build_dir} failed with exit status {returncode}; " + f"see {log_path} for details" + ) + + def custom_function(self, kind: str, particle_count: int, pid_options: list): + """The ms.CustomFunction of [custom_functions] , or None.""" + api = self.custom_functions.get(kind) + if api is None: + return None + function_type = getattr(ms.CustomFunction, kind) + return ms.CustomFunction(api, function_type, particle_count, pid_options) + + def parse_observable(self, name: str, order_observable: str) -> dict: + parts = name.split("-") + sum_momenta = False + sum_observable = False + ordered = False + multiparticles = self.run_card["multiparticles"] + + if len(parts) == 0: + raise ValueError("Invalid observable name") + elif len(parts) == 1: + # event-level observables + obs_name = parts[0] + select_pids = [] + else: + if parts[-1] == "sum": + sum_observable = True + obs_name = parts[-2] + selection = parts[:-2] + elif parts[-2] == "sum": + sum_momenta = True + obs_name = parts[-1] + selection = parts[:-2] + else: + obs_name = parts[-1] + selection = parts[:-1] + select_pids = [] + order_indices = [] + for mp_name in selection: + mp_parts = mp_name.split("_") + if mp_parts[-1].isnumeric(): + order_indices.append(int(mp_parts[-1])) + select_pids.append(multiparticles["_".join(mp_parts[:-1])]) + ordered = True + else: + order_indices.append(0) + select_pids.append(multiparticles[mp_name]) + + return dict( + observable=obs_name, + select_pids=select_pids, + sum_momenta=sum_momenta, + sum_observable=sum_observable, + order_observable=order_observable if ordered else None, + order_indices=order_indices if ordered else [], + ignore_incoming=True, + name=name, + ) + + def init_cuts(self) -> None: + inf = float("inf") + order_observable = self.run_card["cuts"].get("order_by", "pt") + self.cut_data = [ + CutItem( + observable_kwargs=self.parse_observable(key, order_observable), + min=values.get("min", -inf), + max=values.get("max", inf), + mode=values.get("mode", "all"), + ) + for key, values in self.run_card["cuts"].items() + if key != "order_by" + ] + + # [histograms] key that means "the distribution of the event weight", + # normalised to the unit weight sigma_abs, rather than an observable of the + # momenta: an unweighted sample is a spike at 1 (and at -1 for its negative + # weights) and a partially unweighted one shows its spread. + weight_histogram_key = "weight" + + def init_histograms(self) -> None: + order_observable = self.run_card["histograms"].get("order_by", "pt") + self.hist_data = [] + for key, values in self.run_card["histograms"].items(): + if key == "order_by": + continue + from_weight = key == self.weight_histogram_key + self.hist_data.append(HistItem( + observable_kwargs=( + None if from_weight + else self.parse_observable(key, order_observable)), + min=values["min"], + max=values["max"], + bin_count=values["bin_count"], + name=key, + from_weight=from_weight, + )) + + def ensure_pdf_set(self, pdf_set: str) -> misc.LhapdfPaths: + """Make sure the requested LHAPDF set is available, downloading it into + the first writable PDF directory if needed, and return the (possibly + updated) LHAPDF resolution. Nothing is downloaded when the set is + already present.""" + paths = lhapdf_paths() + if paths.find_set(pdf_set): + return paths + if not paths.config: + logger.debug("no usable lhapdf-config: cannot download %s", pdf_set) + return paths + if not paths.download_path: + logger.debug("no writable PDF directory: cannot download %s", pdf_set) + return paths + try: + from madgraph.interface.common_run_interface import CommonRunCmd + os.makedirs(paths.download_path, exist_ok=True) + logger.info("PDF set %s not found; downloading into %s", + pdf_set, paths.download_path) + CommonRunCmd.install_lhapdf_pdfset_static( + paths.config, paths.download_path, pdf_set) + except Exception as err: + logger.warning("Could not download PDF set %s: %s", pdf_set, err) + return paths + options._LHAPDF = paths.with_data_path(paths.download_path) + return options._LHAPDF + + def _no_pdf_message(self, pdf_set: str) -> str: + """Explain a missing PDF set in terms of the settings that fix it.""" + searched = os.pathsep.join(self.lhapdf.data_paths) or \ + "(no PDF set directory could be located)" + if self.lhapdf.config: + found = "using lhapdf-config: %s" % self.lhapdf.config + else: + found = ("no usable lhapdf-config was found, so the set could not " + "be downloaded automatically") + return ( + "LHAPDF set %r, requested by [beam] pdf in Cards/run_card.toml, " + "was not found.\n" + " searched: %s\n" + " %s\n" + "Fix this in one of the following ways:\n" + " * install LHAPDF for this MadGraph: MG7> install lhapdf6\n" + " * point MadGraph at an existing one:\n" + " MG7> set lhapdf /path/to/lhapdf-config\n" + " MG7> save options\n" + " * or set it for this directory only, by adding\n" + " lhapdf = /path/to/lhapdf-config\n" + " to %s\n" + " * or set $LHAPDF_DATA_PATH to a directory containing %r." + % (pdf_set, searched, found, + os.path.join("Cards", "me5_configuration.txt"), pdf_set)) + + def init_beam(self) -> None: + beam_args = self.run_card["beam"] + + # built lazily by build_systematics (needs the subprocess data) + self.systematics = None + self.systematics_data = None + self.systematics_context = None + self.event_histograms = None + self.event_histograms_context = None + + 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, + "transverse_mass": ms.EnergyScale.transverse_mass, + "half_transverse_mass": ms.EnergyScale.half_transverse_mass, + "partonic_energy": ms.EnergyScale.partonic_energy, + } + if beam_args["dynamical_scale_choice"] in dynamical_scales: + dynamical_scale_type = dynamical_scales[beam_args["dynamical_scale_choice"]] + else: + raise ValueError("Unknown dynamical scale choice") + self.scale_kwargs = dict( + dynamical_scale_type=dynamical_scale_type, + ren_scale_fixed=beam_args["fixed_ren_scale"], + fact_scale_fixed=beam_args["fixed_fact_scale"], + ren_scale=beam_args["ren_scale"], + fact_scale1=beam_args["fact_scale1"], + fact_scale2=beam_args["fact_scale2"], + scale_factor=float(beam_args.get("scale_factor", 1.)), + ) + if self.is_decay and self.custom_functions.pop("scale", None) is not None: + logger.warning("custom scale ignored: a decay uses the decaying mass") + if self.is_decay and self.custom_functions.pop("pdf", None) is not None: + logger.warning("custom pdf ignored: a decay has no parton densities") + 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, + scale_factor=1., + ) + self.pdf_grid = None + self.lhapdf = None + self.pdf_dir = None + self.fixed_alphas_info = self.write_fixed_alphas_info() + self.alphas_grid = ms.AlphaSGrid(self.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.pdf_set = pdf_set + self.fixed_alphas_info = None + self.lhapdf = self.ensure_pdf_set(pdf_set) + self.pdf_dir = self.lhapdf.find_set(pdf_set) + if self.pdf_dir is None: + raise RuntimeError(self._no_pdf_message(pdf_set)) + self.alphas_grid = ms.AlphaSGrid(os.path.join(self.pdf_dir, pdf_set, f"{pdf_set}.info")) + if "pdf" in self.custom_functions: + # alpha_s still comes from the PDF set of the run card + self.pdf_grid = ms.CustomPdf(self.custom_functions["pdf"]) + else: + self.pdf_grid = ms.PdfGrid( + os.path.join(self.pdf_dir, pdf_set, f"{pdf_set}_0000.dat") + ) + for context in self.contexts: + if not isinstance(self.pdf_grid, ms.CustomPdf): + self.pdf_grid.initialize_globals(context) + self.alphas_grid.initialize_globals(context) + self.running_coupling = ms.RunningCoupling(self.alphas_grid) + + # ------------------------------------------------------------------ + # scale / PDF systematics + # ------------------------------------------------------------------ + def systematics_enabled(self) -> bool: + """The native scale/PDF weights are requested ([systematics] enable, or + the legacy [generation] systematics alias).""" + try: + enabled = bool(self.run_card["systematics"]["enable"]) + except Exception: + enabled = False + if not enabled: + try: + enabled = bool(self.run_card["generation"]["systematics"]) + except Exception: + enabled = False + if enabled and "pdf" in getattr(self, "custom_functions", {}): + # the variations re-evaluate the PDF set, which a custom PDF replaces + if not getattr(self, "_warned_custom_pdf_systematics", False): + logger.warning("systematics are disabled with a custom pdf") + self._warned_custom_pdf_systematics = True + return False + return enabled + + @staticmethod + def pdf_set_info(info_path: str) -> dict: + """Read SetIndex, NumMembers, ErrorType and SetDesc from an LHAPDF + .info file (missing keys get neutral defaults).""" + info = {"SetIndex": 0, "NumMembers": 1, "ErrorType": "", "SetDesc": ""} + try: + with open(info_path) as f: + for line in f: + key, sep, value = line.partition(":") + key = key.strip() + if not sep or key not in info: + continue + value = value.strip().strip('"').strip("'") + if key in ("SetIndex", "NumMembers"): + info[key] = int(value) + else: + info[key] = value + except OSError as err: + logger.warning("could not read %s: %s", info_path, err) + return info + + def resolve_pdf_set_name(self, entry: str) -> tuple[str, int | None]: + """Turn a [systematics] pdf entry ("", "", "@", + "@") into (set name, member or None) using the LHAPDF + pdfsets.index when an id is given.""" + member = None + if "@" in entry: + entry, member_str = entry.split("@", 1) + member = int(member_str) + entry = entry.strip() + if not entry.isdigit(): + return entry, member + lhaid = int(entry) + index_path = os.path.join(self.pdf_dir or "", "pdfsets.index") + sets = [] + try: + with open(index_path) as f: + for line in f: + parts = line.split() + if len(parts) >= 2 and parts[0].isdigit(): + sets.append((int(parts[0]), parts[1])) + except OSError: + raise ValueError( + "cannot resolve LHAPDF id %s: no %s" % (lhaid, index_path)) + for set_id, name in sets: + if set_id == lhaid: + return name, member + # an id inside a set's member range + sets.sort() + for (set_id, name), (next_id, _) in zip(sets, sets[1:] + [(None, None)]): + if set_id < lhaid and (next_id is None or lhaid < next_id): + return name, lhaid - set_id if member is None else member + raise ValueError("cannot resolve LHAPDF id %s" % lhaid) + + def resolve_pdf_variations(self) -> list: + """Expand the [systematics] pdf entries into ms.PdfMemberSpec objects + (downloading the sets when needed). The nominal member itself is never + listed: the calculator adds it as the central weight of its group.""" + entries = self.run_card["systematics"]["pdf"] + if isinstance(entries, str): + entries = [entries] + specs = [] + seen = set() + for raw in entries: + entry = str(raw).strip() + if not entry or entry.lower() in ("none", "no", "false"): + continue + if entry.lower() == "central": + continue + if entry.lower() == "errorset": + set_name, member = self.pdf_set, None + else: + set_name, member = self.resolve_pdf_set_name(entry) + variation_lhapdf = self.ensure_pdf_set(set_name) + variation_dir = variation_lhapdf.find_set(set_name) + if variation_dir is None: + logger.warning("PDF set %s not available; its variation is skipped", set_name) + continue + info_path = os.path.join(variation_dir, set_name, f"{set_name}.info") + if not os.path.exists(info_path): + logger.warning("PDF set %s not available; its variation is skipped", set_name) + continue + info = self.pdf_set_info(info_path) + is_nominal = set_name == self.pdf_set + if member is not None: + members = [member] + else: + members = list(range(info["NumMembers"])) + if is_nominal: + members = members[1:] + if not members: + logger.warning( + "PDF set %s has a single member: no PDF uncertainty can be " + "computed from it. Use a set with error members in " + "[systematics] pdf (e.g. a PDF4LHC or NNPDF error set).", + set_name) + for m in members: + if is_nominal and m == 0: + continue + if (set_name, m) in seen: + continue + seen.add((set_name, m)) + grid_file = os.path.join(variation_dir, set_name, f"{set_name}_{m:04d}.dat") + if not os.path.exists(grid_file): + logger.warning("PDF member file %s not found; skipped", grid_file) + continue + specs.append(ms.PdfMemberSpec( + set_name=set_name, set_lhaid=info["SetIndex"], member=m, + grid_file=grid_file, info_file=info_path, + error_type=info["ErrorType"], description=info["SetDesc"], + )) + return specs + + def build_systematics_args(self) -> list: + """One ms.SubprocessSystArgs per (unmerged) subprocess: the alpha_s power + of |M|^2 and the beam parton ids of every flavor index, indexed like the + LHECompleter (subprocess_index / flavor_index of the combined events).""" + args = [] + for meta in self.subprocess_data: + beam_pdgs = [ + [int(flavor["options"][0][0]), int(flavor["options"][0][1])] + for flavor in meta["flavors"] + ] + args.append(ms.SubprocessSystArgs( + qcd_power=int(meta.get("qcd_power", -1)), beam_pdgs=beam_pdgs)) + return args + + def build_systematics(self): + """Create the ms.SystematicsCalculator from the [systematics] section, or + None when disabled. Also keeps the JSON description needed by gridpacks.""" + if not self.systematics_enabled(): + self.systematics = None + return None + syst = self.run_card["systematics"] + config = ms.SystematicsConfig() + config.mur = [float(v) for v in syst["mur"]] + config.muf = [float(v) for v in syst["muf"]] + config.together = bool(syst["together"]) + config.dyn_scales = self.resolve_dynamical_scales() + # alternative dynamical scales get the same factor as the nominal one + config.scale_factor = float(self.run_card["beam"].get("scale_factor", 1.)) + config.write_inputs = bool(syst["write_inputs"]) + config.has_pdf = not self.leptonic + # Without parton luminosity -- a decay, or leptonic beams -- there is no + # nominal PDF *set* to describe and no PDF variation to compute, so + # has_pdf gates the nominal-set metadata and the member list below. + # The alpha_s grid is a separate matter: the mu_R variations always need + # one, so its .info file has to be recorded whatever has_pdf says. + # init_beam() skips the PDF lookup for a decay only -- leptonic beams + # still resolve a set and take alpha_s from it -- so the path follows + # self.pdf_dir, and falls back to the constant grid written for a decay. + if self.pdf_dir is not None: + info_path = os.path.join(self.pdf_dir, self.pdf_set, f"{self.pdf_set}.info") + else: + info_path = self.fixed_alphas_info + if config.has_pdf: + info = self.pdf_set_info(info_path) + config.nominal_set_name = self.pdf_set + config.nominal_lhaid = info["SetIndex"] + config.nominal_error_type = info["ErrorType"] + config.nominal_description = info["SetDesc"] + config.pdf_members = self.resolve_pdf_variations() + else: + config.pdf_members = [] + args = self.build_systematics_args() + # the PDFs, alpha_s and (for mixed-order subprocesses) the matrix + # elements are evaluated with the batched madspace functions on this + # CPU context + self.systematics_context = ms.Context(device=ms.cpu_device(), thread_count=1) + matrix_elements, flavor_remap, me_backend = self.build_systematics_matrix_elements( + args, self.systematics_context) + self.systematics = ms.SystematicsCalculator( + config, args, self.pdf_grid, self.alphas_grid, + self.systematics_context, matrix_elements, flavor_remap) + for warning in self.systematics.warnings: + logger.warning("systematics: %s", warning) + self.systematics_data = { + "config": json.loads(config.to_json()), + "subproc_args": [json.loads(a.to_json()) for a in args], + "nominal_grid_file": os.path.join( + self.pdf_dir, self.pdf_set, f"{self.pdf_set}_0000.dat") + if config.has_pdf else None, + "nominal_info_file": info_path, + # matrix elements re-evaluated for the mixed-order subprocesses + "me_paths": [meta["me_path"] for meta in self.subprocess_data], + "me_backend": me_backend, + "flavor_remap": flavor_remap, + "me_frame": self.me_frame, + "incoming_count": self.incoming_count, + } + return self.systematics + + _DYN_SCALE_CODES = { + "transverse_energy": 1, "transverse_mass": 2, + "half_transverse_mass": 3, "partonic_energy": 4, + } + + def resolve_dynamical_scales(self) -> list: + """[systematics] dynamical_scale entries (names or systematics.py codes + 1-4) as codes, without the choice the events were generated with.""" + try: + entries = self.run_card["systematics"]["dynamical_scale"] + except Exception: + return [] + if isinstance(entries, str): + entries = [entries] + beam = self.run_card["beam"] + generated = None + if not (beam["fixed_ren_scale"] and beam["fixed_fact_scale"]): + generated = self._DYN_SCALE_CODES.get(beam["dynamical_scale_choice"]) + codes = [] + for raw in entries: + entry = str(raw).strip().lower() + if not entry: + continue + if entry.isdigit(): + code = int(entry) + elif entry in self._DYN_SCALE_CODES: + code = self._DYN_SCALE_CODES[entry] + else: + raise ValueError("unknown dynamical scale choice %r in [systematics] " + "dynamical_scale" % raw) + if code not in (1, 2, 3, 4): + raise ValueError("dynamical scale code %s must be 1-4" % code) + if code == generated or code in codes: + continue + codes.append(code) + return codes + + def build_systematics_matrix_elements(self, args, context): + """For the subprocesses whose |M|^2 mixes several powers of alpha_s + (qcd_power = -1), load the matrix element into the systematics CPU + context so that the renormalisation scale variations can re-evaluate it + at the varied alpha_s. Returns ([MatrixElement or None per subprocess], + flavor remap, backend name); empty when nothing is needed or no CPU + library exists.""" + need = [i for i, a in enumerate(args) if a.qcd_power < 0] + empty = ([], [], None) + if not need or self.run_card["run"]["dummy_matrix_element"]: + return empty + cpu_backends = [b for b in self.backends if not b.startswith(("cuda", "hip"))] + if not cpu_backends: + logger.warning("systematics: no CPU matrix-element library available, the " + "renormalisation scale variations of the mixed-order " + "subprocesses cannot be computed") + return empty + backend = cpu_backends[0] + matrix_elements = [] + flavor_remap = [] + for i, meta in enumerate(self.subprocess_data): + flavor_remap.append([int(f["index"]) for f in meta["flavors"]]) + if i not in need: + matrix_elements.append(None) + continue + api = context.load_matrix_element( + meta["me_path"].format(device=backend), self.param_card_path, + me_parameters(self.run_card)) + matrix_elements.append(ms.MatrixElement( + api, + [ms.MatrixElement.momenta_in, ms.MatrixElement.alpha_s_in, + ms.MatrixElement.flavor_in], + [ms.MatrixElement.matrix_element_out], + False, + self.me_frame, + self.incoming_count, + )) + logger.info("systematics: re-evaluating the matrix element of %d mixed-order " + "subprocess(es) for the renormalisation scale variations", len(need)) + return matrix_elements, flavor_remap, backend + + def build_event_histograms(self): + """ms.EventHistograms for the [histograms] observables, filled with the + written events and all their variation weights (info.json + "event_histograms"); None without histograms.""" + if not self.hist_data \ + or not self.run_card["run"]["postprocessing_histograms"]: + self.event_histograms = None + return None + context = ms.Context(device=ms.cpu_device(), thread_count=1) + specs = [ + ms.EventHistogramSpec( + name=item.name, min=item.min, max=item.max, + bin_count=item.bin_count, from_weight=item.from_weight) + for item in self.hist_data + ] + from_momenta = [item for item in self.hist_data if not item.from_weight] + observables = [] + for meta in self.subprocess_data: + if not from_momenta: + # nothing to evaluate on the momenta (only the weight is + # histogrammed): no observable runtime to build + observables.append(None) + continue + all_pids = clean_pids(meta["incoming"]) + clean_pids(meta["outgoing"]) + values = ms.ObservableValues([ + build_observable( + observable_pids(meta, item.observable_kwargs), + len(meta["incoming"]), item.observable_kwargs + ) + for item in from_momenta + ]) + observables.append(ms.SubprocessObservables(values, len(all_pids))) + # the weight histograms are drawn in units of the unit weight sigma_abs + # (the cross section of a sample without negative weights) + self.event_histograms_context = context + self.event_histograms = ms.EventHistograms( + context, specs, observables, + reference_weight=self._mean_event_weight()) + return self.event_histograms + + def _mean_event_weight(self) -> float: + """The unit weight of the unweighted events, sigma_abs, used as the + unit of the weight histograms: the cross section of a positive sample, + and for one with negative weights (an interference) what puts them at + +-1 rather than dividing by a signed, possibly small, total. 0 (the raw + weight) when it is not available.""" + try: + return float(self.event_generator.status().mean_abs) + except Exception as error: + logger.debug("no cross section for the weight histograms: %s", error) + return 0. + + def log_systematics_summary(self) -> None: + """Print the scale/PDF uncertainties on the total cross section (the + per-variation cross sections are in events.weights.json / info.json).""" + if self.systematics is None or self.systematics.weight_count == 0: + return + summary = json.loads(self.systematics.summary()) + # the percentages come from systematics_summary, which get_result() + # reads too: the box and the scan row must not disagree, so neither + # computes them itself + xsec = systematics_summary.nominal_cross_section(summary) + if xsec is None: + return + def format_variation(up, down): + # right-justify the signed numbers (not the sign alone) so the % + # values line up across rows, with extra spacing between the two + return f"{'+%.3g' % up:>6}% {'-%.3g' % down:>6}%" + + scale_count = sum(1 for v in self.systematics.variations if v.is_scale) + member_count = len(self.systematics.members) + rows = [("Variations per event:", + f"{self.systematics.weight_count} ({scale_count} scale, {member_count} PDF members)")] + for pdf in summary.get("pdf", []): + rows.append(("PDF set:", f"{pdf['pdf_set']}, {pdf['error_type']}")) + rows.append(("Original cross-section:", f"{xsec} pb")) + scale = systematics_summary.scale_percentages(summary) + if scale is not None: + rows.append(("Scale variation:", format_variation(*scale))) + for _entry, up, down in systematics_summary.pdf_percentages(summary): + rows.append(("PDF variation:", format_variation(up, down))) + if self.event_generator_config.verbosity == ms.Verbosity.pretty: + box = ms.PrettyBox("Systematics", len(rows), [24, 0]) + box.set_column(0, [label for label, _ in rows]) + box.set_column(1, [value for _, value in rows]) + box.print_first() + else: + for label, value in rows: + logger.info("systematics, %s %s", label, value) + + 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 needs_systematics_matrix_elements(self) -> bool: + """The systematics will re-evaluate a matrix element per scale point. + + True when the native weights are on and some subprocess mixes alpha_s + powers (qcd_power < 0), i.e. when build_systematics_matrix_elements() + loads a matrix element. Conservative: it does not check that a CPU + backend exists, so it can be True where that call ends up loading + nothing. + """ + if not self.systematics_enabled(): + return False + if self.run_card["run"]["dummy_matrix_element"]: + return False + return any(int(meta.get("qcd_power", -1)) < 0 + for meta in self.subprocess_data) + + def init_generator_config(self) -> None: + run_args = self.run_card["run"] + gen_args = self.run_card["generation"] + vegas_args = self.run_card["vegas"] + cfg = ms.GeneratorConfig() + cfg.target_count = gen_args["events"] + cfg.vegas_damping = vegas_args["damping"] + cfg.max_overweight_truncation = gen_args["max_overweight_truncation"] + cfg.freeze_max_weight_after = gen_args["freeze_max_weight_after"] + first_pass = None + if self.surrogate is not None and self.surrogate["mode"] == "unweight": + first_pass = self.surrogate + amplicol = self.run_card["amplicol"] + if ( + first_pass is None and amplicol["mode"] == "unweight" + and amplicol["color_accuracy"] != "lc" + and any(is_amplicol(meta) for meta in self.subprocess_data) + ): + first_pass = amplicol + if first_pass is not None: + # [generation] applies to the second pass, which gives the events + cfg.max_overweight_truncation = ( + first_pass["first_pass_max_overweight_truncation"] + ) + cfg.freeze_max_weight_after = ( + first_pass["first_pass_freeze_max_weight_after"] + ) + cfg.start_batch_size = vegas_args["start_batch_size"] + cfg.max_batch_size = vegas_args["max_batch_size"] + cfg.survey_min_iters = gen_args["survey_min_iters"] + cfg.survey_max_iters = gen_args["survey_max_iters"] + cfg.survey_target_precision = gen_args["survey_target_precision"] + cfg.optimization_patience = vegas_args["optimization_patience"] + cfg.optimization_threshold = vegas_args["optimization_threshold"] + cfg.cpu_batch_size = gen_args["cpu_batch_size"] + cfg.gpu_batch_size = gen_args["gpu_batch_size"] + cfg.verbosity = resolve_verbosity(run_args["verbosity"]) + cfg.combine_thread_count = run_args["combine_thread_pool_size"] + if self.needs_systematics_matrix_elements() and cfg.combine_thread_count != 1: + # SystematicsCalculator::compute() runs inside the combine workers, + # and a matrix element's per-thread process instances live in a + # ThreadResource indexed by ThreadPool::thread_index() -- the + # *calling* thread's index. A parallel combine pool would therefore + # need one process instance per worker; building N of them to index + # them safely is not worth it, the more so as + # SystematicsCalculator::matrix_elements() already serialises every + # evaluation on its own mutex, so that parallelism was not being + # used anyway. Run the combine stage single-threaded instead. + logger.info( + "systematics: the combine stage runs single-threaded because a " + "subprocess needs its matrix element re-evaluated for the " + "renormalisation scale variations") + cfg.combine_thread_count = 1 + cfg.cut_efficiency_threshold = gen_args["cut_efficiency_threshold"] + cfg.max_cut_repetitions = gen_args["max_cut_repetitions"] + self.event_generator_config = cfg + self.event_generator = None + + def init_subprocesses(self) -> None: + self.backends = self.compile_matrix_elements() + self.subprocesses = [] + if self.merged_subprocess_data is None: + for subproc_id, meta in enumerate(self.subprocess_data): + subproc_class = ( + AmplicolSubprocess if is_amplicol(meta) else MadgraphSubprocess + ) + self.subprocesses.append(subproc_class(self, meta, subproc_id)) + else: + for subproc_id, meta in enumerate(self.merged_subprocess_data): + # an AmpliCol subprocess is never merged with another one + first = self.subprocess_data[meta["subprocesses"][0]] + subproc_class = ( + AmplicolSubprocess if is_amplicol(first) else MadgraphSubprocess + ) + self.subprocesses.append( + subproc_class(self, meta, subproc_id, self.subprocess_data) + ) + + def compile_matrix_elements(self) -> list[str]: + """Build the matrix-element library of every subprocess, and return the + build backend of each requested device (in the same order). + + A 'cpu' device builds the backend named by the 'cpu_mode' entry; the + cuda/hip devices ignore cpu_mode and build their own backend. A + cpu_mode of 'auto' is resolved to the actual SIMD backend chosen on + this machine. + + SubProcesses/makefile is a dispatcher over the P* directories, so a + single 'make -j N' there builds all the subprocesses at once with one + shared pool of N jobs: no subprocess is built with N jobs while the + others wait, and none is limited to N/#subprocesses jobs either. + """ + device_names = self.run_card["run"]["device"] + if not isinstance(device_names, list): + device_names = [device_names] + cpu_mode = self.run_card["run"]["cpu_mode"] + backends = [backend_of(name, cpu_mode) for name in device_names] + self.compile_amplicol() + madmatrix_data = [ + meta for meta in self.subprocess_data if not is_amplicol(meta) + ] + if not madmatrix_data: + return backends + + first_proc_path = madmatrix_data[0]["path"] + subproc_path = os.path.dirname(first_proc_path) + + fptype = fptype_of(self.run_card["run"]["precision"]) + + # Resolve cpu_mode='auto' once (the build rules pick the best SIMD + # backend available here), so all subprocesses agree on the library names. + auto_backend = None + if "auto" in backends: + auto_backend = resolve_auto_backend(first_proc_path) + logger.info("cpu_mode 'auto' resolved as '%s'", auto_backend) + resolved = [ + auto_backend if backend == "auto" else backend + for backend in backends + ] + + nb_core = self.run_card["run"]["cpu_thread_pool_size"] + if not nb_core or nb_core < 0: + nb_core = os.cpu_count() or 1 + + log_path = os.path.join(self.run_path, "compile_subprocesses.log") + for backend in resolved: + missing = [ + meta for meta in madmatrix_data + if not os.path.isfile(meta["me_path"].format(device=backend)) + ] + if not missing: + continue + logger.info( + f"Start compilation of SubProcesses for device '{backend}' " + f"({len(missing)} subprocess(es), {nb_core} parallel job(s)), " + f"see log detail in {log_path}" + ) + start_time = time.time() + self.make_subprocesses( + subproc_path, + [f"BACKEND={backend}", f"FPTYPE={fptype}", "USEBUILDDIR=1"], + nb_core, log_path, + ) + logger.info( + f"Compilation of SubProcesses done in {time.time() - start_time:.1f} s" + ) + return resolved + + def compile_amplicol(self) -> None: + """Build the libraries of the [amplicol] processes that are missing.""" + paths = sorted({ + meta["path"] for meta in self.subprocess_data + if is_amplicol(meta) and not ( + os.path.isfile(meta["me_path"]) and os.path.isfile(meta["ratio_path"]) + ) + }) + log_path = os.path.join(self.run_path, "compile_amplicol.log") + for path in paths: + logger.info(f"Compiling the AmpliCol matrix elements in {path}, see {log_path}") + self.make_subprocesses(path, [], 1, log_path) + + @staticmethod + def make_subprocesses( + subproc_path: str, args: list[str], nb_core: int, log_path: str + ) -> None: + """Run 'make -j nb_core' in SubProcesses/, appending the full build + output to log_path (one run per device, so the file is appended to).""" + command = ["make", f"-j{nb_core}"] + args + with open(log_path, "a") as log: + log.write(f"\n$ cd {subproc_path} && {' '.join(command)}\n") + log.flush() + returncode = subprocess.call( + command, cwd=subproc_path, stdout=log, stderr=subprocess.STDOUT + ) + if returncode != 0: + raise RuntimeError( + f"Compilation of the SubProcesses failed with exit status " + f"{returncode}; see {log_path} for details" + ) + + def build_event_generator(self, phasespaces: list[PhaseSpace]) -> ms.EventGenerator: + channel_generators = [] + subprocess_channels = [[] for _ in self.subprocesses] + for i, (subproc, phasespace) in enumerate(zip(self.subprocesses, phasespaces)): + for integrand, channel in zip( + subproc.build_integrands(phasespace), phasespace.channels + ): + if channel.event_generator is None: + channel.event_generator = ms.ChannelEventGenerator( + contexts=self.contexts, + integrand=integrand, + event_file=os.path.join(self.temp_path, f"events.{i}.{channel.name}.npy"), + weight_file=os.path.join(self.temp_path, f"weights.{i}.{channel.name}.npy"), + config=self.event_generator_config, + subprocess_index=i, + name=f"{i}.{channel.name}", + histograms=subproc.histograms + ) + subprocess_channels[i].append(len(channel_generators)) + channel_generators.append(channel.event_generator) + + event_generator = ms.EventGenerator( + contexts=self.contexts, + channels=channel_generators, + status_file=self.status_file, + config=self.event_generator_config, + seed=self.run_seed, + groups=self.weighting_groups(subprocess_channels), + ) + unused_globals = ( + set(self.contexts[0].global_names()) - event_generator.used_globals() + ) + for context in self.contexts: + for global_name in unused_globals: + context.delete_global(global_name) + return event_generator + + def survey_phasespaces( + self, phasespaces: list[PhaseSpace | None], survey_pass: int = 0 + ) -> ms.EventGenerator | None: + ps_filtered = [ps for ps in phasespaces if ps is not None] + if len(ps_filtered) == 0: + return None + event_generator = self.build_event_generator(ps_filtered) + event_generator.survey(survey_pass) + return event_generator + + def survey(self) -> None: + # survey_pass distinguishes the survey() calls below: "both" mode can + # re-survey a channel carried over unchanged from the multichannel pass + # into the final (simplified) pass, and both passes schedule jobs on the + # same underlying ChannelEventGenerator. The explicit pass index keeps each + # pass's job seeds independent of the other passes' job counts, rather than + # depending on call history. + 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 -- and 'both'/'auto' end up + # simplifying towards it. 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 in ["multichannel", "both", "auto"]: + self.phasespaces = [ + subproc.build_multichannel_phasespace() + for subproc in self.subprocesses + ] + self.event_generator = self.survey_phasespaces(self.phasespaces, 0) + elif phasespace_mode == "flat": + self.phasespaces = [ + subproc.build_flat_phasespace() + for subproc in self.subprocesses + ] + self.event_generator = self.survey_phasespaces(self.phasespaces, 0) + else: + raise ValueError("Unknown phasespace mode") + + channel_status = self.event_generator.channel_status() + chan_offset = 0 + madnis_enabled = False + for subproc, ps in zip(self.subprocesses, self.phasespaces): + mean = 0. + variance = 0. + count_opt = 0 + for status in channel_status[chan_offset:chan_offset + len(ps.channels)]: + mean += status.mean_abs + variance += status.error_abs**2 + count_opt += status.count_opt + # A subprocess whose matrix element is identically zero (an FCNC + # channel with every Wilson coefficient at zero, say) integrates to + # exactly zero, so there is no relative spread to speak of. Report 0 + # rather than dividing by it: that asks for the smallest networks and + # leaves MadNIS off, which is what an empty subprocess wants. + rsd = (variance * count_opt)**0.5 / mean if mean else 0. + subproc.set_madnis_auto_settings(rsd) + chan_offset += len(ps.channels) + if subproc.madnis_settings["enable"]: + madnis_enabled = True + if not ( + phasespace_mode == "both" or (phasespace_mode == "auto" and madnis_enabled) + ): + return + + phasespaces_multi = self.phasespaces + cross_sections = [] + index = 0 + for phasespace in phasespaces_multi: + channel_count = len(phasespace.channels) + # rank the channels by the integral of |w|: a channel whose signed + # integral cancels (an interference) can still carry most of the + # variance, and folding it into the flat channel would blow it up + cross_sections.append([ + status.mean_abs + for status in channel_status[index:index + channel_count] + ]) + index += channel_count + + self.phasespaces = [ + subproc.simplify_phasespace(ps_multi, cross_secs) + for subproc, ps_multi, cross_secs in zip( + self.subprocesses, phasespaces_multi, cross_sections + ) + ] + if any( + ps_multi is not ps_both + for ps_multi, ps_both in zip(phasespaces_multi, self.phasespaces) + ): + # distinct survey_pass: a channel carried over unchanged from the + # multichannel pass (pass 0) into this resurvey must not share its + # seed stream with that earlier pass. + self.event_generator = self.survey_phasespaces(self.phasespaces, 1) + + def train_madnis(self) -> None: + madnis_args = self.run_card["madnis"] + if not any(subproc.madnis_settings["enable"] for subproc in self.subprocesses): + return + + gen_args = self.run_card["generation"] + run_args = self.run_card["run"] + + verbosity = resolve_verbosity(run_args["verbosity"]) + madnis_phasespaces = [] + training_args = [] + self.event_generator = None + for subproc, phasespace in zip(self.subprocesses, self.phasespaces): + for channel in phasespace.channels: + channel.event_generator = None + + config = ms.MadnisConfig() + config.learning_rate = subproc.madnis_settings["lr"] + config.batches = subproc.madnis_settings["train_batches"] + config.log_interval = madnis_args["log_interval"] + config.integration_history_length = madnis_args["integration_history_length"] + config.channel_dropping_interval = madnis_args["channel_dropping_interval"] + config.channel_dropping_threshold = madnis_args["channel_dropping_threshold"] + config.cpu_generator_batch_size = gen_args["cpu_batch_size"] + config.gpu_generator_batch_size = gen_args["gpu_batch_size"] + config.gpu_generator_batch_granularity = madnis_args["gpu_generator_batch_granularity"] + config.generator_target_size_factor = madnis_args["generator_target_size_factor"] + config.batch_size_offset = madnis_args["batch_size_offset"] + config.batch_size_per_channel = subproc.madnis_settings["batch_size_per_channel"] + config.uniform_channel_ratio = madnis_args["uniform_channel_ratio"] + config.lr_schedule = madnis_args["lr_scheduler"] + config.adam_beta1 = madnis_args["adam_beta1"] + config.adam_beta2 = madnis_args["adam_beta2"] + config.adam_eps = madnis_args["adam_eps"] + config.adam_weight_decay = madnis_args["adam_weight_decay"] + config.grad_clip_threshold = madnis_args["grad_clip_threshold"] + config.buffer_capacity = madnis_args["buffer_capacity"] + config.minimum_buffer_size = madnis_args["minimum_buffer_size"] + config.buffered_steps_fraction = madnis_args["buffered_steps_fraction"] + config.buffer_skip_batches = madnis_args["buffer_skip_batches"] + config.buffer_unweighting_quantile = madnis_args["buffer_unweighting_quantile"] + config.fixed_cwnet_fraction = subproc.madnis_settings["fixed_cwnet_fraction"] + config.softclip_threshold = madnis_args["softclip_threshold"] + config.compressed_channel_weight_count = madnis_args["compressed_channel_weight_count"] + phasespace = subproc.build_madnis(phasespace) + madnis_phasespaces.append(phasespace) + training_args.append( + ms.TrainingArgs( + config=config, + integrands=subproc.build_integrands( + phasespace, + madnis_training=True, + drop_cuts_and_rescale=madnis_args["drop_zero_integrands"] + ), + cwnet=phasespace.cwnet, + ) + ) + + gen_context = self.contexts[0] + opt_context = ms.Context( + device=self.devices[0], thread_count=self.pool_sizes[0] + ) + opt_context.copy_globals_from(gen_context) + + madnis_training = ms.MultiMadnisTraining( + generator_context=gen_context, + optimizer_context=opt_context, + training_args=training_args, + verbosity=verbosity, + status_file=self.status_file, + # Reuses the run's resolved seed (also used by build_event_generator()). + # Only the single-channel CPU sample-generation path is currently seeded + # -- buffered training and GPU multi-channel batches are still + # non-deterministic. + seed=self.run_seed, + ) + madnis_training.train() + for phasespace, active_channels in zip( + madnis_phasespaces, madnis_training.active_channels() + ): + phasespace.channels = [ + phasespace.channels[index] for index in active_channels + ] + del madnis_training + del opt_context + self.phasespaces = madnis_phasespaces + for context in self.contexts[1:]: + context.copy_globals_from(self.contexts[0]) + self.event_generator = self.build_event_generator(madnis_phasespaces) + self.madnis_trained = True + + def update_madnis_status_single( + self, batch: int, batch_target: int, loss: float, lr: float, channel_count: int + ) -> None: + now = time.time() + if batch + 1 < batch_target: + if now - self.last_update_time < 0.1: + return + self.last_update_time = now + progress_bar = ms.format_progress((batch + 1) / batch_target, 52) + time_diff = now - self.madnis_wall_time + time_str = f"{format_time(time_diff)}" + else: + progress_bar = "" + wall_diff = now - self.madnis_wall_time + cpu_diff = time.process_time() - self.madnis_cpu_time + time_str = ( + f"{format_time(wall_diff, centi=True)} wall, " + f"{format_time(cpu_diff, centi=True)} cpu" + ) + batch_str = f"{batch + 1} / {batch_target}" + self.madnis_box.set_column(1, [ + f"{batch_str:<15} {progress_bar}", + f"{loss:>.4f}", + f"{channel_count}", + time_str + ]) + self.madnis_box.print_update() + + def update_madnis_status_multi( + self, + subproc_id: int, + batch: int, + batch_target: int, + loss: float, + lr: float, + channel_count: int + ) -> None: + now = time.time() + subproc_count = len(self.subprocesses) + if batch + 1 < batch_target: + if now - self.last_update_time < 0.1: + return + self.last_update_time = now + progress_bar = ms.format_progress((batch + 1) / batch_target, 34) + progress_bar_all = ms.format_progress( + (subproc_id * batch_target + batch + 1) / (subproc_count * batch_target), + 52 + ) + time_str = f"{format_time(now - self.madnis_wall_time)}" + subproc_str = f"{subproc_id} / {subproc_count}" + elif subproc_id < subproc_count - 1: + progress_bar = "" + progress_bar_all = ms.format_progress( + ((subproc_id + 1) * batch_target + 1) / (subproc_count * batch_target), + 52 + ) + time_str = f"{format_time(now - self.madnis_wall_time)}" + subproc_str = f"{subproc_id} / {subproc_count}" + else: + progress_bar = "" + progress_bar_all = "" + wall_diff = now - self.madnis_wall_time + cpu_diff = time.process_time() - self.madnis_cpu_time + time_str = ( + f"{format_time(wall_diff, centi=True)} wall, " + f"{format_time(cpu_diff, centi=True)} cpu" + ) + subproc_str = f"{subproc_count} / {subproc_count}" + batch_str = f"{batch + 1} / {batch_target}" + self.madnis_upper_box.set_column(1, [ + f"{subproc_str:<15} {progress_bar_all}", + time_str, + ]) + self.madnis_lower_box.set_row(subproc_id + 1, [ + f"{subproc_id}", + f"{loss:>.4f}", + f"{channel_count}", + f"{batch_str:<15} {progress_bar}", + ]) + self.madnis_upper_box.print_update() + self.madnis_lower_box.print_update() + + @property + def gridpack_run_mode(self) -> str: + """How this run ends: a gridpack run may skip the events (see the + [gridpack] run_mode run card option); any other run is a regular one.""" + if not self.run_card["gridpack"]["save_gridpack"]: + return "regular" + return self.run_card["gridpack"]["run_mode"] + + def finish_without_events(self) -> None: + """The end of a gridpack run that writes no events: only what the + gridpack needs is prepared, then it is stored.""" + if self.gridpack_run_mode == "minimal": + # a trained MadNIS needs nothing more, a VEGAS grid has to converge + if not self.madnis_trained: + self.event_generator.optimize() + else: + self.event_generator.fix_max_weights() + self.build_systematics() + self.save_gridpack() + + def generate_events(self) -> None: + if self.gridpack_run_mode != "regular": + self.finish_without_events() + return + start_time = get_start_time() + self.event_generator.generate() + output_format = self.run_card["run"]["output_format"] + systematics = self.build_systematics() + histograms = self.build_event_histograms() + if output_format == "compact_npy": + self.lhe_completer = None + self.event_generator.combine_to_compact_npy( + os.path.join(self.run_path, "events.npy"), systematics, histograms + ) + self.write_lhe_header(systematics) + self.save_lhe_completer() + elif output_format == "lhe_npy": + self.lhe_completer = self.build_lhe_completer() + self.event_generator.combine_to_lhe_npy( + os.path.join(self.run_path, "events.npy"), self.lhe_completer, + systematics, histograms + ) + self.write_lhe_header(systematics) + elif output_format == "lhe": + self.lhe_completer = self.build_lhe_completer() + lhe_path = os.path.join(self.run_path, "events.lhe") + self.event_generator.combine_to_lhe( + lhe_path, self.lhe_completer, + self.build_lhe_meta(), systematics, histograms + ) + # Ship the LHE compressed by default. These files are large and + # very compressible, madevent has always stored its events + # gzipped, and every consumer here already accepts either form + # (see _find_event_file). misc.gzip replaces events.lhe with + # events.lhe.gz, and switches to an external multithreaded tool + # above 256 MB. + misc.gzip(lhe_path) + else: + raise ValueError("Unknown output format") + self.write_hwu(histograms) + self.make_plots(histograms) + if systematics is not None: + self.write_systematics_sidecar() + self.log_systematics_summary() + self.save_gridpack() + + def save_lhe_completer(self) -> None: + """Keep the LHE completer next to compact_npy events, so npy_to_lhe + can complete them into an LHE file later (a shower or MadSpin requested + after the run). Never fatal to the run.""" + try: + self.build_lhe_completer().save( + os.path.join(self.run_path, npy_to_lhe.COMPLETER_FILE)) + except Exception as err: + logger.warning("could not save the LHE completer; these events " + "cannot be converted to LHE later: %s", err) + + hwu_file_name = "MADatLO.HwU" + plot_dir_name = "plots" + + def make_plots(self, histograms) -> None: + """Draw the event-sample histograms into Events//plots, with the + scale and PDF bands info.json carries, and say where they landed. + + This does not go through the HwU file: the bands are already computed + per bin by madspace, so drawing them directly is both fewer steps and + more than histograms.py can do without the LHAPDF python module. + """ + if histograms is None or not self.run_card["run"]["make_plots"]: + return + out_dir = os.path.join(self.run_path, self.plot_dir_name) + try: + data = json.loads(histograms.to_json(self.systematics)) + written = plots.render(data, out_dir) + except plots.BackendMissing as error: + logger.info("matplotlib is not available (%s): no plots. The " + "distributions are in %s", + error, os.path.join(self.run_path, "info.json")) + return + except Exception as error: + logger.warning("could not draw the histograms: %s", error) + return + if not written: + return + logger.info("%d plot(s) stored in %s", len(written), out_dir) + + def write_hwu(self, histograms) -> None: + """Write the event-sample histograms next to the events in the HwU + format (MADatLO.HwU), the one an aMC@NLO run writes as MADatNLO.HwU + and madgraph/various/histograms.py reads: an mg7 distribution can then + be overlaid on an NLO one without converting anything in between. The + numbers are the ones info.json carries, bands included, so this is off + by default -- it is a second copy in another format, for the runs that + are going to be plotted.""" + if histograms is None or not self.run_card["run"]["write_hwu"]: + return + try: + data = json.loads(histograms.to_json(self.systematics)) + summary = (json.loads(self.systematics.summary()) + if self.systematics is not None else None) + text = hwu_output.to_hwu(data, summary) + except Exception as error: + logger.warning("could not write the HwU histograms: %s", error) + return + if not text: + return + with open(os.path.join(self.run_path, self.hwu_file_name), "w") as f: + f.write(text) + + def write_systematics_sidecar(self) -> None: + """Describe the variation weights next to the event file + (events.weights.json): ids, scale factors, PDF set/member and the + text, so npy consumers get the same metadata as the LHE + header, plus the per-variation cross sections.""" + if self.systematics is None: + return + data = json.loads(self.systematics.summary()) + data["initrwgt"] = self.systematics.initrwgt() + data["columns"] = ["rwgt_%d" % i for i in self.systematics.weight_ids] + with open(os.path.join(self.run_path, "events.weights.json"), "w") as f: + json.dump(data, f, indent=1) + + @staticmethod + def _histogram_mean(hist_min, hist_max, bin_values): + """Cross-section-weighted mean of a histogrammed observable over its + range. ``bin_values`` carries the underflow and overflow bins first and + last, as both madspace histogram kinds do; they have no position and + are left out.""" + values = list(bin_values)[1:-1] + total = sum(values) + if not values or total == 0: + return None + width = (hist_max - hist_min) / len(values) + return sum(v * (hist_min + (i + 0.5) * width) + for i, v in enumerate(values)) / total + + def _histogram_means(self) -> dict: + """name -> mean of every [histograms] entry, from the post-processing + histograms (the final event sample) when they were filled, otherwise + from the weighted integration histograms.""" + if self.event_histograms is not None: + data = json.loads(self.event_histograms.to_json(self.systematics)) + hists = [(h["name"], h["min"], h["max"], h["bin_values"]) + for h in data] + else: + hists = [(h.name, h.min, h.max, h.bin_values) + for h in self.event_generator.histograms()] + means = {} + for name, hist_min, hist_max, values in hists: + mean = self._histogram_mean(hist_min, hist_max, values) + if mean is not None: + means[name] = mean + return means + + def get_result(self) -> dict: + """Return the run result: cross-section (pb) with MC error, the number + of (unweighted) events, and the mean of every observable declared in the + [histograms] section. Used to build the scan summary.""" + status = self.event_generator.status() + result = {'cross(pb)': status.mean, 'error(pb)': status.error, + 'nb_event': status.count_unweighted} + try: + if self.systematics is not None and self.systematics.weight_count: + summary = json.loads(self.systematics.summary()) + scale = systematics_summary.scale_percentages(summary) + if scale is not None: + result['scale_up(%)'], result['scale_down(%)'] = scale + pdf = systematics_summary.pdf_percentages(summary) + if pdf: + # one column pair; the nominal set's entry comes first + _entry, up, down = pdf[0] + result['pdf_up(%)'], result['pdf_down(%)'] = up, down + except Exception as err: + logger.warning("could not extract the systematics summary: %s", err) + try: + for name, mean in self._histogram_means().items(): + result['<%s>' % name] = mean + except Exception as err: + logger.warning("could not extract observable means: %s", err) + return result + + def _beam_info(self): + """Return (beam_pdg_ids, beam_energies) for the LHE block. + + `incoming` holds the *partonic* initial state (e.g. gluons), not the + 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] + data = self.subprocess_data[0] + incoming = data["incoming"] + flavor0 = data["flavors"][0]["options"][0] + init_pdgs = flavor0[:len(incoming)] + beam_pdgs = [pdg if abs(code) in (81, 82) else code + for code, pdg in zip(incoming, init_pdgs)] + return beam_pdgs, [half_e, half_e] + + def _lhapdf_id(self): + """Central LHAPDF id of the beam PDF set (read from its .info SetIndex), + or -1 for a leptonic beam (no PDF).""" + if self.leptonic or not self.pdf_dir: + return -1 + pdf_set = self.run_card["beam"]["pdf"] + info = os.path.join(self.pdf_dir, pdf_set, "%s.info" % pdf_set) + try: + for line in open(info): + if line.strip().startswith("SetIndex:"): + return int(line.split(":", 1)[1].strip()) + except Exception as err: + logger.warning("could not read LHAPDF id from %s: %s", info, err) + return -1 + + def build_lhe_meta(self, systematics=None): + """Build the LHE header/ metadata: the param_card () and the + run_card.toml () headers plus the beam/PDF/cross-section info + needed by downstream tools (systematics, MadSpin, ...). + + `combine_to_lhe` injects the header itself, so callers that + go through it should leave `systematics` unset; pass it only when + building meta for a writer that bypasses that injection (e.g. + write_lhe_header).""" + beam_pdgs, energies = self._beam_info() + lhaid = self._lhapdf_id() + pdf_group = -1 if self.leptonic else 0 + status = self.event_generator.status() + xsec, err = status.mean, status.error + weight_mode, max_weight = lhe_weight_info(status) + with open(self.param_card_path) as f: + param_text = f.read() + with open(os.path.join("Cards", "run_card.toml")) as f: + run_text = f.read() + headers = [] + # generation commands (model + process): read by MadSpin/reweight/... + proc_card = os.path.join("Cards", "proc_card_mg5.dat") + if os.path.exists(proc_card): + with open(proc_card) as f: + headers.append(ms.LHEHeader(name="MG5ProcCard", content=f.read())) + headers.append(ms.LHEHeader(name="slha", content=param_text)) + headers.append(ms.LHEHeader(name="MG7RunCard", content=run_text)) + # The resolved seed (even when the run_card requested a random one via + # seed = -1), so the run can be reproduced from the LHE file alone. + headers.append(ms.LHEHeader(name="MG7Seed", content=str(self.run_seed))) + if systematics is not None and systematics.weight_ids: + headers.append(ms.LHEHeader( + name="initrwgt", content=systematics.initrwgt(), escape_content=False + )) + return ms.LHEMeta( + beam1_pdg_id=beam_pdgs[0], beam2_pdg_id=beam_pdgs[1], + beam1_energy=energies[0], beam2_energy=energies[1], + beam1_pdf_authors=pdf_group, beam2_pdf_authors=pdf_group, + beam1_pdf_id=lhaid, beam2_pdf_id=lhaid, + weight_mode=weight_mode, + # positional: the pybind arg name for max_weight is non-kwarg-safe + processes=[ms.LHEProcess(xsec, err, max_weight, 1)], + headers=headers, + ) + + def write_lhe_header(self, systematics) -> None: + """Write header.lhe next to events.npy: the
/ blocks + (run card, param card, beam/PDF, cross section) that the npy formats + otherwise drop, with no events.""" + header_path = os.path.join(self.run_path, "header.lhe") + writer = ms.LHEFileWriter(header_path, self.build_lhe_meta(systematics)) + del writer # closes the file (writes the closing tag) + + def build_lhe_completer(self): + all_mcdata = [ + None if is_amplicol(meta) + else self.subprocesses[index].build_multi_channel_data() + if self.merged_subprocess_data is None + else build_multi_channel_data(meta, self) + for index, meta in enumerate(self.subprocess_data) + ] + subproc_args = [ + amplicol_subproc_args(meta, self) if mcdata is None else ms.SubprocArgs( + topologies = [topo[0] for topo in mcdata.topologies], + permutations = mcdata.permutations, + diagram_indices = mcdata.diagram_indices, + diagram_color_indices = mcdata.diagram_color_indices, + diagram_propagator_pdgs = mcdata.diagram_propagator_pdgs, + color_flows = meta["color_flows"], + pdg_color_types = { + int(key): value + for key, value in meta["pdg_color_types"].items() + }, + helicities = lhe_helicities(meta, self.run_card), + pdg_ids = [flavor["options"] for flavor in meta["flavors"]], + ) + for mcdata, meta in zip(all_mcdata, self.subprocess_data) + ] + return ms.LHECompleter( + subproc_args=subproc_args, + bw_cutoff=self.run_card["phasespace"]["bw_cutoff"] + ) + + def save_gridpack(self) -> None: + if not self.run_card["gridpack"]["save_gridpack"]: + return + + gridpack_path = os.path.join(self.run_path, "gridpack") + data_path = os.path.join(gridpack_path, "data") + events_path = os.path.join(gridpack_path, "Events") + os.mkdir(gridpack_path) + os.mkdir(data_path) + os.mkdir(events_path) + self.contexts[0].save_globals(os.path.join(data_path, "globals.tar")) + + # one json object per channel, joined by hand so that the channel + # json is not parsed and dumped again + channel_entries = [ + f"{json.dumps(channel.status().name)}:{channel.to_json(self.gridpack_run_mode == "fix_max_weight")}" + for channel in self.event_generator.channels() + ] + with open(os.path.join(data_path, "channels.json"), "w") as f: + f.write("{" + ",".join(channel_entries) + "}") + + lib_path = os.path.join(gridpack_path, "lib") + if self.run_card["gridpack"]["include_source"]: + os.mkdir(lib_path) + shutil.copytree("src", os.path.join(gridpack_path, "src")) + shutil.copytree("SubProcesses", os.path.join(gridpack_path, "SubProcesses")) + shutil.copytree("backend", os.path.join(gridpack_path, "backend")) + else: + shutil.copytree("lib", lib_path) + for _, path in self.custom_function_paths: + target = path if os.path.isabs(path) else os.path.join(gridpack_path, path) + if not os.path.exists(target): + os.makedirs(os.path.dirname(target), exist_ok=True) + shutil.copy(path, target) + + if self.run_card["gridpack"]["include_madspace_source"]: + shutil.copytree( + _MADSPACE_DIR, + os.path.join(gridpack_path, "madspace"), + ignore=shutil.ignore_patterns("build", "install"), + ) + + if self.run_card["gridpack"]["include_madspace"]: + shutil.copytree( + _INSTALL_DIR / "madspace", + os.path.join(gridpack_path, "madspace", "install", "madspace"), + ignore=shutil.ignore_patterns("include"), + ) + + matrix_elements = [] + for subproc in self.subprocess_data: + me_path = subproc["me_path"] + matrix_elements.append(me_path) + + cards_path = os.path.join(gridpack_path, "Cards") + os.mkdir(cards_path) + shutil.copy(os.path.join("Cards", "param_card.dat"), cards_path) + # Full run card with a header noting it is read-only in gridpack context. + import io as _io + _buf = _io.StringIO() + self.run_card.write(_buf) + _header = ( + "# This is the run card used to generate this gridpack.\n" + "# Modifying this file will have no effect on gridpack execution.\n" + "# To change event-generation settings, edit grid_run_card.toml.\n\n" + ) + with open(os.path.join(cards_path, "run_card.toml"), 'w') as _f: + _f.write(_header + _buf.getvalue()) + # Minimal card containing only the settings used by generate_events. + # The gridpack ships the libraries that were actually built, and their + # names carry the resolved backend, so the card it runs from has to name + # that backend as well: a cpu_mode of 'auto' would send the gridpack + # looking for a ..._auto.so that was never built. The full run_card.toml + # written just above keeps 'auto', since that is what was asked for. + device_names = self.run_card["run"]["device"] + if not isinstance(device_names, list): + device_names = [device_names] + resolved_cpu_mode = None + for name, backend in zip(device_names, getattr(self, "backends", [])): + if name.split(":")[0] == "cpu": + resolved_cpu_mode = backend + break + previous_cpu_mode = self.run_card["run"]["cpu_mode"] + if resolved_cpu_mode is not None: + self.run_card["run"]["cpu_mode"] = resolved_cpu_mode + # A gridpack is run to feed a shower or detector chain, which reads + # LHE, so it writes LHE by default whatever format this run used + # (bin/generate_events --output_format still selects an npy output). + previous_output_format = self.run_card["run"]["output_format"] + self.run_card["run"]["output_format"] = "lhe" + # the directories of this run mean nothing on the machine running the gridpack + previous_dirs = (self.run_card["gridpack"]["temp_output_dir"], + self.run_card["gridpack"]["output_dir"]) + self.run_card["gridpack"]["temp_output_dir"] = "" + self.run_card["gridpack"]["output_dir"] = "" + try: + self.run_card.write_gridpack_card( + os.path.join(cards_path, "grid_run_card.toml")) + finally: + self.run_card["run"]["cpu_mode"] = previous_cpu_mode + self.run_card["run"]["output_format"] = previous_output_format + (self.run_card["gridpack"]["temp_output_dir"], + self.run_card["gridpack"]["output_dir"]) = previous_dirs + + bin_path = os.path.join(gridpack_path, "bin") + os.mkdir(bin_path) + gen_events_file = os.path.join(bin_path, "generate_events") + template_dir = os.path.dirname(__file__) + shutil.copy(os.path.join(template_dir, "gridpack.py"), gen_events_file) + shutil.copy(os.path.join(template_dir, "custom_function_loader.py"), bin_path) + os.chmod(gen_events_file, 0o755) + shutil.copy(os.path.join(template_dir, "gridpack_setup.py"), bin_path) + # streaming needs the fixed maximum weights of this mode + if self.gridpack_run_mode == "fix_max_weight": + shutil.copy(os.path.join(template_dir, "event_stream.py"), bin_path) + + data = { + "matrix_elements": matrix_elements, + # loaded before the matrix elements, as in the run + "custom_functions": self.custom_function_paths, + "me_parameters": me_parameters(self.run_card), + "source_hash": ms.SOURCE_HASH, + } + with open(os.path.join(data_path, "data.json"), "w") as f: + json.dump(data, f) + + if self.lhe_completer is None: + self.lhe_completer = self.build_lhe_completer() + self.lhe_completer.save(os.path.join(data_path, "lhe.json")) + self.save_gridpack_lhe_meta(os.path.join(data_path, "lhe_meta.json")) + if self.systematics_data is not None: + systematics_data = dict(self.systematics_data) + # A decay takes alpha_s from a constant grid written into the run + # directory rather than from an LHAPDF set, so the gridpack cannot + # look that file up: ship it and point at the shipped copy. + if self.fixed_alphas_info and \ + systematics_data["nominal_info_file"] == self.fixed_alphas_info: + name = os.path.basename(self.fixed_alphas_info) + shutil.copy(self.fixed_alphas_info, os.path.join(data_path, name)) + systematics_data["nominal_info_file"] = os.path.join("data", name) + with open(os.path.join(data_path, "systematics.json"), "w") as f: + json.dump(systematics_data, f) + if self.run_card["gridpack"]["compress"]: + shutil.make_archive(gridpack_path, "gztar", self.run_path, "gridpack") + shutil.rmtree(gridpack_path) + + def save_gridpack_lhe_meta(self, path) -> None: + """The run-independent part of the LHE header (cards, beams, PDF), for + the gridpack to write complete LHE headers: it adds the cross section + and seed of its own run, which is why those two are left out here.""" + meta = self.build_lhe_meta() + data = {key: getattr(meta, key) for key in ( + "beam1_pdg_id", "beam2_pdg_id", "beam1_energy", "beam2_energy", + "beam1_pdf_authors", "beam2_pdf_authors", "beam1_pdf_id", + "beam2_pdf_id", "weight_mode")} + data["headers"] = [ + {"name": h.name, "content": h.content, + "escape_content": h.escape_content} + for h in meta.headers if h.name != "MG7Seed"] + with open(path, "w") as f: + json.dump(data, f) + + def get_mass(self, pid: int) -> float: + return self.param_card.get_value("mass", pid) + + def get_width(self, pid: int) -> float: + return self.param_card.get_value("width", pid) diff --git a/madgraph/iolibs/template_files/mg7/model_widths.py b/madgraph/iolibs/template_files/mg7/model_widths.py new file mode 100644 index 0000000000..f007473a2a --- /dev/null +++ b/madgraph/iolibs/template_files/mg7/model_widths.py @@ -0,0 +1,155 @@ +"""The model of a process directory and its auto widths. Does not import madspace.""" + +import logging +import os +import re +import subprocess + +from madgraph.various import misc +from madgraph.iolibs.template_files.mg7.bootstrap import MG_ROOT as _MG_ROOT + +logger = logging.getLogger("madgraph7") + + +_model_hash_warned = set() +_model_cache = {} + + +def read_stored_model(me_dir=None) -> "str | None": + """The model this process was generated with, as recorded at output time in + ``SubProcesses/model.txt``: a UFO path carrying its restriction (e.g. + ``.../models/sm-no_b_mass``) or a model name -- in either case something + ``import model`` understands. Returns None when the process was written + without that information. + + Warns (once per directory) when the model on disk no longer matches the + hash stored on the second line, i.e. it changed since the output was made. + """ + base = me_dir or os.getcwd() + model_file = os.path.join(base, "SubProcesses", "model.txt") + if not os.path.exists(model_file): + return None + with open(model_file) as f: + lines = f.read().splitlines() + model = lines[0].strip() if lines else "" + stored_hash = lines[1].strip() if len(lines) > 1 else "" + if not model: + logger.warning("SubProcesses/model.txt is empty; the model of this " + "process is unknown.") + return None + + # verify the model on disk still matches the one used at output time. The + # reference may carry a restriction suffix ('sm-no_b_mass'); the hash was + # taken on the UFO directory itself ('sm'). + model_dir = model if os.path.isdir(model) else model.rsplit('-', 1)[0] + if stored_hash and os.path.isdir(model_dir) \ + and model_file not in _model_hash_warned: + current_hash = misc.hash_model_files(model_dir) + if current_hash and current_hash != stored_hash: + _model_hash_warned.add(model_file) + logger.warning( + "The model at %s has changed since this process was generated " + "(hash mismatch); anything computed from it now (the 'auto' " + "widths, the dependent masses/widths of the param_card) may be " + "inconsistent with the matrix element.", model_dir) + return model + + +def _proc_characteristic(me_dir=None) -> dict: + """SubProcesses/proc_characteristics as a plain dict (empty when the file is + missing or unreadable). Plain, because ProcCharacteristic is a ConfigFile, + whose get() is the parameter accessor and takes no default.""" + from madgraph.various import banner as _banner_mod + path = os.path.join(me_dir or os.getcwd(), 'SubProcesses', + 'proc_characteristics') + if os.path.exists(path): + try: + return dict(_banner_mod.ProcCharacteristic(path)) + except Exception as error: + logger.debug("could not read %s: %s", path, error) + return {} + + +def load_process_model(me_dir=None): + """Import the UFO model this process was generated with, as + :func:`models.import_ufo.import_model` builds it -- the object every + model-level consumer of the param_card expects (see + :meth:`MG7Cmd.get_model`). None when the process recorded no model, or when + that model no longer imports. + + Cached per (reference, complex-mass-scheme): importing a model is slow and + the card question asks for it again at every 'update dependent'.""" + ref = read_stored_model(me_dir) + if not ref: + return None + cms = _proc_characteristic(me_dir).get('complex_mass_scheme', False) + key = (ref, bool(cms)) + if key not in _model_cache: + import models.import_ufo as import_ufo + try: + with misc.MuteLogger(['madgraph.model'], [50]): + _model_cache[key] = import_ufo.import_model( + ref, complex_mass_scheme=bool(cms)) + except Exception as error: + # 'update dependent' calls this under a SIGALRM whose handler + # raises a TimeOutError of its own (a class defined inside the + # caller, so it can only be recognised by name). Let it through: + # the caller then says the model took too long to load and how to + # force it, instead of reporting a model that does not import. + if type(error).__name__ == 'TimeOutError': + raise + logger.debug("could not import the model %s: %s", ref, error) + _model_cache[key] = None + return _model_cache[key] + + +def compute_auto_widths(param_card_path=os.path.join("Cards", "param_card.dat")) -> None: + """Fill any width set to ``auto`` in the param_card, using madgraph and the + model stored at output time (``SubProcesses/model.txt``), and write the + result back into the card. A no-op when the card has no ``auto`` width. + + Called before every generation, so a single run and each scan point (whose + param_card is rewritten just before ``run_single``) both get model-computed + widths.""" + try: + with open(param_card_path) as f: + text = f.read() + except OSError: + return + # matches "DECAY auto" (optionally "auto@NLO"), as in + # common_run_interface.static_check_param_card + pdgs = re.findall(r"(?im)^\s*decay\s+([+-]?\d+)\s+auto", text) + if not pdgs: + return + pdgs = list(dict.fromkeys(pdgs)) # de-duplicate, keep order + + model = read_stored_model() + if not model: + logger.warning( + "The param_card requests 'auto' width(s) for %s but the model of " + "this process is unknown; leaving them as-is.", " ".join(pdgs)) + return + + mg5 = str(_MG_ROOT / "bin" / "madgraph") + if not os.path.exists(mg5): + logger.warning("Cannot find madgraph at %s; 'auto' widths not computed.", mg5) + return + + import tempfile + cmds = "import model %s\ncompute_widths %s --path=%s\n" % ( + model, " ".join(pdgs), os.path.abspath(param_card_path)) + with tempfile.NamedTemporaryFile("w", suffix=".mg5", delete=False) as fh: + fh.write(cmds) + cmdfile = fh.name + try: + logger.info("Computing 'auto' width(s) for %s ...", " ".join(pdgs)) + proc = subprocess.run([mg5, cmdfile]) + if proc.returncode != 0: + logger.warning( + "compute_widths returned a non-zero exit code; the param_card " + "may still contain 'auto' entries.") + finally: + try: + os.remove(cmdfile) + except OSError: + pass diff --git a/madgraph/iolibs/template_files/mg7/observables.py b/madgraph/iolibs/template_files/mg7/observables.py new file mode 100644 index 0000000000..27523e7d73 --- /dev/null +++ b/madgraph/iolibs/template_files/mg7/observables.py @@ -0,0 +1,178 @@ +"""Observables, cuts and histograms of the run card, and PDG id helpers.""" + +from dataclasses import dataclass + +from madgraph.iolibs.template_files.mg7.common import ms + +@dataclass +class CutItem: + observable_kwargs: dict + min: float + max: float + mode: str + + +@dataclass +class HistItem: + observable_kwargs: dict + min: float + max: float + bin_count: int + # the distribution of the event weight itself (the reserved "weight" key + # of [histograms]); it is read off the event record, not computed from + # the momenta, so it has no observable + name: str = "" + from_weight: bool = False + + +# 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 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 + + +# observables that tell a particle from its antiparticle, and so need the signed +# pdg ids of the actual flavours rather than clean_pids' unsigned representatives +_SIGNED_OBSERVABLES = {"sfos_pair_mass"} + + +def _flavor_options(meta: dict, unmerged_meta: dict | None = None) -> list[list[int]]: + """Every flavour assignment (signed pdg ids, incoming then outgoing) a + subprocess evaluates with the same momenta.""" + if "flavors" in meta: + return [ + option for flavor in meta["flavors"] for option in flavor["options"] + ] + return [ + option + for subproc in meta.get("subprocesses", []) + for option in _flavor_options(unmerged_meta[subproc]) + ] + + +def observable_pids( + meta: dict, observable_kwargs: dict, unmerged_meta: dict | None = None +) -> list[int]: + """The pdg ids an observable of this subprocess is built with. + + Most observables only select particles by group, which the unsigned + representatives of clean_pids do. One that pairs a particle with its + antiparticle needs the real signed flavours: those of the first flavour + assignment, provided every assignment of the subprocess forms the same + pairs -- the cut is applied once to momenta that all of them share. + """ + all_pids = clean_pids(meta["incoming"]) + clean_pids(meta["outgoing"]) + if observable_kwargs.get("observable") not in _SIGNED_OBSERVABLES: + return all_pids + options = _flavor_options(meta, unmerged_meta) + if not options: + raise ValueError( + f"{observable_kwargs.get('name')}: no flavour information for " + f"subprocess {meta['incoming']} > {meta['outgoing']}" + ) + n_in = len(meta["incoming"]) + groups = [set(group) for group in observable_kwargs["select_pids"]] + if len(groups) == 1: + groups = groups * 2 + + def sfos_pairs(pids): + out = range(n_in, len(pids)) + return frozenset( + frozenset((i, j)) + for i in out if pids[i] in groups[0] + for j in out if pids[j] in groups[1] + if i != j and pids[i] != 0 and pids[i] == -pids[j] + ) + + reference = sfos_pairs(options[0]) + if any(sfos_pairs(option) != reference for option in options[1:]): + raise ValueError( + f"{observable_kwargs.get('name')}: the flavours grouped into subprocess " + f"{meta['incoming']} > {meta['outgoing']} pair the particles " + "differently, so a same-flavour opposite-sign cut cannot be applied " + "to their common momenta; regenerate the process after " + "'set apply_flavor_grouping False'" + ) + return list(options[0]) + + +def build_observable(pids: list[int], n_in: int, observable_kwargs: dict): + """ms.Observable of a subprocess with `n_in` incoming particles. + + madspace takes the first two momenta to be the beams (ignore_incoming), so + for a 1 -> N decay it would drop the first decay product from every + selection. There the selection runs over all the momenta instead, with the + incoming particle set to pdg id 0, which is in no group. + """ + if n_in == 2: + return ms.Observable(pids, **observable_kwargs) + pids = [0] * n_in + list(pids[n_in:]) + return ms.Observable(pids, **dict(observable_kwargs, ignore_incoming=False)) + + +def decay_products(meta: dict, unmerged_meta: dict | None = None) -> set[int]: + """Outgoing positions (0-based, beams excluded) of the particles that come + from an on-shell (decay-chain) propagator, as MadEvent's check_decay marks + them. Every channel of a subprocess shares its decay chains, so the first + channel tells, read through its first diagram's permutation.""" + if not meta["channels"]: + return set() + channel = meta["channels"][0] + topo_channel = channel + if "vertices" not in channel: + topo_channel = unmerged_meta[channel["subprocess"]]["channels"][channel["channel"]] + on_shell = topo_channel.get("on_shell_propagators", []) + if not on_shell: + return set() + vertices = topo_channel["vertices"] + # the closing vertex creates no propagator + created = {vertex[-1]: vertex[:-1] for vertex in vertices[:-1]} + + def legs(name): + if name in created: + return set().union(*(legs(child) for child in created[name])) + return {name} + + n_in = len(meta["incoming"]) + n_out = len(meta["outgoing"]) + products = set() + for prop in on_shell: + external = legs(f"p{prop}") + incoming = {leg for leg in external if leg.startswith("i")} + outgoing = {int(leg[1:]) for leg in external if leg.startswith("o")} + if not incoming: + products |= outgoing + elif len(incoming) == n_in: + # built from the beams: the decay is everything else + products |= set(range(n_out)) - outgoing + # event[i] = topology[permutation[i]] + permutation = channel["diagrams"][0]["permutation"] + return { + i - n_in for i in range(n_in, len(permutation)) + if permutation[i] - n_in in products + } + + +def pid_is_qcd(pid: int): + return abs(pid) in [21, 1, 2, 3, 4, 5, 6, 81] diff --git a/madgraph/iolibs/template_files/mg7/options.py b/madgraph/iolibs/template_files/mg7/options.py new file mode 100644 index 0000000000..9225c37271 --- /dev/null +++ b/madgraph/iolibs/template_files/mg7/options.py @@ -0,0 +1,86 @@ +"""The MG5 options of the process directory and the LHAPDF paths they select.""" + +import os + +from madgraph.various import misc + +_ROOTED_OPTIONS = ('lhapdf', 'lhapdf_py3', 'lhapdf_py2', 'heptools_install_dir') + +def load_mg5_options(me_dir=None) -> dict: + """Read the tool paths from the MadGraph configuration, so the launcher + knows where LHAPDF and the optional programs (Pythia8/Delphes/MadSpin/ + reweight/analysis) live. + + Files are read least- to most-specific, each overriding the previous, which + is the same layering CommonRunCmd.set_configuration uses: the card inside + this process directory has the last word, then this installation, then the + per-user file. Relative values are resolved against the root of the file + they came from. + + ``me_dir`` is the process directory whose Cards/me5_configuration.txt has + the last word. It defaults to the working directory, which is right for + bin/generate_events (it chdirs there first) but not for an in-process + caller, which has not moved yet. + """ + + import madgraph + mg5dir = os.path.dirname(os.path.dirname(os.path.abspath(madgraph.__file__))) + me_dir = os.getcwd() if me_dir is None else me_dir + + options = { + 'pythia-pgs_path': None, 'pythia8_path': None, + 'madanalysis5_path': None, 'exrootanalysis_path': None, 'delphes_path': None, + 'rivet_path': None, 'contur_path': None, 'f2py_compiler': None, + 'lhapdf': None, 'lhapdf_py3': None, 'lhapdf_py2': None, 'timeout': 0, + 'heptools_install_dir': None, 'hepmc3_path': None, + } + config_files = [] + base_config = misc.base_config_file() + if base_config: + config_files.append((base_config, mg5dir)) + user_config = misc.user_config_file() + if user_config: + config_files.append((user_config, mg5dir)) + config_files.append((misc.install_config_file(mg5dir), mg5dir)) + config_files.append((os.path.join(me_dir, 'Cards', 'me5_configuration.txt'), + me_dir)) + for cfg, root in config_files: + if not os.path.exists(cfg): + continue + with open(cfg) as fsock: + for line in fsock: + line = line.split('#', 1)[0] + if '=' not in line: + continue + name, value = (x.strip() for x in line.split('=', 1)) + if name not in options or value in ('', 'None'): + continue + if (name.endswith('_path') or name in _ROOTED_OPTIONS) and \ + not os.path.isabs(value) and \ + (os.sep in value or value.startswith('.')): + value = os.path.normpath(os.path.join(root, value)) + options[name] = value + options['mg5_path'] = mg5dir # enables MadSpin/reweight + return options + + +_LHAPDF = None + +def lhapdf_paths(refresh: bool = False) -> misc.LhapdfPaths: + """The LHAPDF installation this run should use, resolved once. + + Uses the same helper as the 'launch' command, so driving a process + directory straight through bin/generate_events behaves identically to + launching it from MadGraph. + """ + + global _LHAPDF + if _LHAPDF is None or refresh: + _LHAPDF = misc.resolve_lhapdf(load_mg5_options()) + # Export what we found: the real LHAPDF library, systematics and the + # MadSpin/reweight subprocesses read these from the environment. + if _LHAPDF.data_paths: + os.environ["LHAPDF_DATA_PATH"] = os.pathsep.join(_LHAPDF.data_paths) + if _LHAPDF.config: + os.environ["MADGRAPH_LHAPDF_CONFIG"] = _LHAPDF.config + return _LHAPDF diff --git a/madgraph/iolibs/template_files/mg7/postprocess.py b/madgraph/iolibs/template_files/mg7/postprocess.py new file mode 100644 index 0000000000..38bb38cf76 --- /dev/null +++ b/madgraph/iolibs/template_files/mg7/postprocess.py @@ -0,0 +1,380 @@ +"""Post-processing tools run on the events, and their logging.""" + +import logging +import os +import sys +import time + +from madgraph.iolibs.template_files.mg7 import npy_to_lhe +from madgraph.iolibs.template_files.mg7.options import lhapdf_paths, load_mg5_options + +def _find_event_file(run_path): + """Locate the LHE event file produced by generate_events (if any).""" + for name in ("events.lhe", "events.lhe.gz"): + path = os.path.join(run_path, name) + if os.path.exists(path): + return path + return None + + +def _find_or_convert_event_file(run_path, log): + """The run's LHE event file, converting its npy events when that is all + the run wrote (see npy_to_lhe).""" + lhe_path = _find_event_file(run_path) + if lhe_path is None and npy_to_lhe.can_convert(run_path): + log.info("No LHE event file in %s: converting the npy events.", run_path) + try: + lhe_path = npy_to_lhe.convert(run_path) + except Exception as error: + _report_failure(log, "npy->LHE conversion", error, run_path) + return lhe_path + + +def _report_failure(log, what, error, directory=None): + """Log a post-processing failure and write the full traceback to a file + (whose path is printed) so the problem can be investigated.""" + import traceback + if not directory or not os.path.isdir(directory): + directory = os.getcwd() + path = os.path.join(directory, "%s_crash.log" % what.replace(" ", "_")) + try: + with open(path, "w") as fsock: + fsock.write(traceback.format_exc()) + except Exception: + path = None + log.warning("%s failed: %s", what, error) + if path: + log.warning("full traceback written to: %s", path) + + +def _off(value) -> bool: + """True when a switch value means the tool was not selected.""" + return value in (None, "OFF", "Not Avail.", "Not Avail. (numpy missing)") + + +def _has_effective_handler(lg) -> bool: + """True when ``lg`` (or an ancestor it propagates to) already has a handler, + i.e. some other party is already displaying these records.""" + while lg: + if lg.handlers: + return True + if not lg.propagate: + return False + lg = lg.parent + return False + + +_TOOL_LOGGING_READY = False + + +def _setup_logging(): + """Make the reused madevent tool drivers' progress visible on screen. + + The drivers already narrate what they do through logger.info / + update_status ("Running Pythia8 [arXiv:...]", "Splitting .lhe event + file...", the live Idle/Running/Completed job counters, "Running + MadSpin", ...), but the standalone mg7 launcher never attaches a handler + to the 'madgraph'/'madevent' loggers, so all of it is swallowed. Attach a + colored INFO console handler to them (idempotent).""" + global _TOOL_LOGGING_READY + if _TOOL_LOGGING_READY: + return + if _has_effective_handler(logging.getLogger("madgraph")): + # Someone already configured logging for us -- in practice MG5, which + # now runs the launch in its own process. Its configuration is + # authoritative (it owns stdout_level and the tutorial logger); adding + # a second handler here would print every line twice. + _TOOL_LOGGING_READY = True + return + try: + import madgraph.interface.coloring_logging # registers ColorFormatter + formatter = logging.ColorFormatter("%(message)s") + except Exception: + formatter = logging.Formatter("%(message)s") + handler = logging.StreamHandler(sys.stdout) + handler.setFormatter(formatter) + handler._mg7_tool_handler = True + for name in ("madgraph", "madevent", "cmdprint", "madgraph7"): + lg = logging.getLogger(name) + if not any(getattr(h, "_mg7_tool_handler", False) for h in lg.handlers): + lg.addHandler(handler) + lg.setLevel(logging.INFO) + lg.propagate = False + _TOOL_LOGGING_READY = True + + +def selected_tools(switch) -> list: + """The post-processing drivers run_selected_tools runs for this switch.""" + # MadAnalysis5 hadron level analyses the shower/detector output, so it only + # makes sense when a shower ran (mirrors madevent's card gating: + # analysis == 'MadAnalysis5' and shower != 'OFF'). + ma5 = switch.get("analysis") == "MadAnalysis5" + showered = not _off(switch.get("shower")) + return [t for t, on in ( + ("reweighting", not _off(switch.get("reweight"))), + ("MadSpin", not _off(switch.get("madspin"))), + ("MadAnalysis5 (parton level)", ma5), + ("Pythia8 shower", switch.get("shower") == "Pythia8"), + ("Delphes", switch.get("detector") == "Delphes"), + ("MadAnalysis5 (hadron level)", ma5 and showered), + ("Rivet", switch.get("analysis") == "Rivet"), + ) if on] + + +def run_selected_tools(switch, process) -> None: + """Run the optional post-processing programs selected in the merged + question on the generated events. + + All the tools are driven through :class:`MG7RunCmd`, a thin adapter around + the standard madevent run interface (``MadEventCmd``). This means the mg7 + output reuses the *exact* madevent tool drivers -- in particular Pythia8 is + showered in parallel over the cluster/multicore backend, identically to a + madevent run -- instead of re-implementing each tool here. The command + sequence mirrors ``MadEventCmd.do_launch`` (reweight, MadSpin, MA5 parton, + shower, Delphes, MA5 hadron, Rivet); each tool checks its own card, so a + tool is invoked only when its switch is on. + """ + log = logging.getLogger("madevent") + + tools = selected_tools(switch) + if not tools: + # A switch that is not off does not necessarily select a driver here: + # "Not Avail." is not off, analysis = ExRoot has no mg7 driver, and a + # shower switch set to something other than Pythia8 selects nothing + # either. Without this a plain generate/output/launch announced a + # post-processing step with nothing after the colon, and paid for + # building the run interface (or an npy->LHE conversion) to do nothing. + return + + lhe_path = _find_or_convert_event_file(process.run_path, log) + if lhe_path is None: + log.warning("No LHE event file in %s; cannot run %s.", + process.run_path, ", ".join(tools)) + return + + run_name = os.path.basename(os.path.dirname(os.path.abspath(lhe_path))) + run_dir = os.path.dirname(os.path.abspath(lhe_path)) + + ma5 = switch.get("analysis") == "MadAnalysis5" + showered = not _off(switch.get("shower")) + + log.info("") + log.info("Post-processing the generated events with: %s", ", ".join(tools)) + + try: + from madgraph.iolibs.template_files.mg7.run_interface import MG7RunCmd + cmd = MG7RunCmd(os.getcwd(), load_mg5_options(), run_name, lhe_path) + cmd.set_run_name(run_name, None, "parton") + except Exception as error: + _report_failure(log, "post-processing setup", error, run_dir) + return + + def run(tool, command): + bar = "=" * 60 + log.info("") + log.info(bar) + log.info(" post-processing step: %s", tool) + log.info(bar) + start = time.time() + try: + cmd.exec_cmd(command, postcmd=False, printcmd=False) + except Exception as error: + _report_failure(log, tool, error, run_dir) + else: + log.info(" -> %s done (%.1fs)", tool, time.time() - start) + + # Same order as MadEventCmd.do_launch. The commands take no run name and so + # act on cmd.run_name: reweight runs on the generated events, then + # decay_events (MadSpin) repoints cmd.run_name to the "_decayed_i" run, + # so the shower and the analyses that follow act on the decayed events -- + # exactly as a madevent run does. Each driver still guards on its own card. + if not _off(switch.get("reweight")): + run("reweight", "reweight -from_cards") + if not _off(switch.get("madspin")): + run("MadSpin", "decay_events -from_cards") + if ma5: + run("MadAnalysis5 (parton)", "madanalysis5_parton --no_default") + if switch.get("shower") == "Pythia8": + run("Pythia8 shower", "shower --no_default") + if switch.get("detector") == "Delphes": + run("Delphes", "delphes --no_default") + # hadron-level MA5 needs the shower/detector output: only run it if a shower + # was requested (otherwise there is no hadron-level event file to analyse). + if ma5 and showered: + run("MadAnalysis5 (hadron)", "madanalysis5_hadron --no_default") + if switch.get("analysis") == "Rivet": + run("Rivet", "rivet --no_default") + + # Finalize the run exactly as MadEventCmd.do_launch does after the + # shower/detector/analysis tools: store_result() processes the deferred + # "to_store" actions -- in particular it gzips the Pythia8 HepMC to + # _pythia8_events.hepmc.gz (the pythia8_card default is HEPMCoutput:file + # = hepmc.gz). The deferred Rivet job reads exactly that .gz path, so this + # must run before the Rivet/Contur postprocessor below. + try: + cmd.store_result() + except Exception as error: + _report_failure(log, "store_result", error, run_dir) + + if switch.get("analysis") == "Rivet": + # do_rivet only *prepared* the run (wrote Events//run_rivet.sh) and + # deferred execution to the postprocessor -- the rivet_card default has + # run_rivet_later = True, so with --no_default it logged "Skipping Rivet + # for now, passing it to postprocessor" and appended the run to + # cmd.postprocessing_dirs. MadEventCmd.do_launch runs that postprocessor + # at the end of a run; do the same here so the .yoda (and, when + # run_contur = True, the Contur limits) are actually produced. + bar = "=" * 60 + log.info("") + log.info(bar) + log.info(" post-processing step: Rivet/Contur postprocessing") + log.info(bar) + start = time.time() + try: + cmd.postprocessing() + except Exception as error: + _report_failure(log, "Rivet-Contur postprocessing", error, run_dir) + else: + log.info(" -> Rivet/Contur postprocessing done (%.1fs)", + time.time() - start) + + +def _add_time_of_flight(lhe_path, threshold, param_card_path, log): + """Add invariant-lifetime (vtim) information to the LHE events, drawing each + unstable particle's decay length from its width (in mm). Mirrors + common_run_interface.CommonRunCmd.do_add_time_of_flight but reads the widths + from the local param_card instead of the LHE banner.""" + import random + try: + import madgraph.various.lhe_parser as lhe_parser + except ImportError: + import internal.lhe_parser as lhe_parser + from models import check_param_card as param_card_mod + from madgraph.various import misc as _misc + import madgraph.iolibs.files as files + + need_zip = lhe_path.endswith('.gz') + if need_zip: + _misc.gunzip(lhe_path) + lhe_path = lhe_path[:-3] + + param_card = param_card_mod.ParamCard(param_card_path) + cst = 6.58211915e-25 # hbar in GeV s + c = 299792458000 # speed of light in mm/s + log.info('Adding time of flight information on %s', lhe_path) + lhe = lhe_parser.EventFile(lhe_path) + out = open('%s_2vertex.lhe' % lhe_path, 'w') + out.write(lhe.banner) + for event in lhe: + for particle in event: + # default=0 -> particles without a decay entry are treated as stable + width = param_card['decay'].get((abs(particle.pid),), 0.).value + if width: + vtim = c * random.expovariate(width / cst) + if vtim > threshold: + particle.vtim = vtim + out.write(str(event)) + out.write('\n') + out.close() + lhe.close() + files.mv('%s_2vertex.lhe' % lhe_path, lhe_path) + if need_zip: + _misc.gzip(lhe_path) + + +def _lhapdf_config_path(): + """Path to lhapdf-config so systematics can import the python lhapdf + module (required to compute PDF/scale variations).""" + return lhapdf_paths().config + + +def _run_systematics(lhe_path, cfg, log): + """Run systematics.py (scale/PDF variations) on the LHE file, reusing the + exact worker madevent uses (systematics.call_systematics).""" + try: + import madgraph.various.systematics as systematics + except ImportError: + import internal.systematics as systematics + + def fmt(vals): + return ','.join(str(v) for v in vals) + + opts = [] + if cfg.get('systematics_mur'): + opts.append('--mur=%s' % fmt(cfg['systematics_mur'])) + if cfg.get('systematics_muf'): + opts.append('--muf=%s' % fmt(cfg['systematics_muf'])) + if cfg.get('systematics_pdf'): + opts.append('--pdf=%s' % fmt(cfg['systematics_pdf'])) + extra = cfg.get('systematics_str_options', '') + if extra: + opts.extend(extra.split()) + # The mg7 LHE has no block, so systematics cannot read the per-event + # LO reweighting info. For a process with a single alpha_s power (stored in + # SubProcesses/proc_characteristics at output time) it can be reconstructed + # from the events: pass that power as --lo_nqcd. -1 means the QCD power is + # not uniform, so the reconstruction is not applicable. + if not any(o.startswith('--lo_nqcd') for o in opts): + try: + from madgraph.various import banner as _banner_mod + except ImportError: + import internal.banner as _banner_mod + pc_path = os.path.join('SubProcesses', 'proc_characteristics') + if os.path.exists(pc_path): + try: + nqcd = int(_banner_mod.ProcCharacteristic(pc_path)['single_qcd_order']) + except Exception: + nqcd = -1 + if nqcd >= 0: + opts.append('--lo_nqcd=%d' % nqcd) + # tell systematics where to find lhapdf (so it can link the python module) + if not any(o.startswith('--lhapdf_config') for o in opts): + lhapdf_config = _lhapdf_config_path() + if lhapdf_config: + opts.append('--lhapdf_config=%s' % lhapdf_config) + + log.info('Running systematics on %s %s', lhe_path, ' '.join(opts)) + systematics.call_systematics([lhe_path, lhe_path] + opts, + log=lambda x: log.info(str(x))) + + +def run_lhe_postprocessing(process) -> None: + """Run the LHE-level post-processings configured in the run_card + [postprocessing] section on the generated event file (displaced-vertex + time-of-flight and systematics). Only applies when an LHE file exists.""" + lhe_path = _find_event_file(process.run_path) + if lhe_path is None: + return + # process.run_card is a RunCardMG7; the section view exposes .get(key, def) + try: + cfg = process.run_card["postprocessing"] + except Exception: + return + log = logging.getLogger('madevent') + + 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'): + native = getattr(process, 'systematics', None) + if native is not None and native.weight_count > 0: + log.warning("[postprocessing] systematics is ignored: the scale/PDF " + "weights were already computed by madspace ([systematics] " + "section); set systematics.enable = false to use the legacy " + "systematics.py path instead.") + else: + try: + _run_systematics(lhe_path, cfg, log) + except Exception as error: + _report_failure(log, "systematics computation", error, + os.path.dirname(lhe_path)) + + tof = cfg.get('time_of_flight', -1.0) + try: + if tof is not None and float(tof) >= 0: + _add_time_of_flight(lhe_path, float(tof), process.param_card_path, log) + except Exception as error: + _report_failure(log, "add_time_of_flight", error, + os.path.dirname(lhe_path)) diff --git a/madgraph/iolibs/template_files/mg7/run_card.toml b/madgraph/iolibs/template_files/mg7/run_card.toml index ea618c4862..0837f4258e 100644 --- a/madgraph/iolibs/template_files/mg7/run_card.toml +++ b/madgraph/iolibs/template_files/mg7/run_card.toml @@ -142,6 +142,15 @@ first_pass_max_overweight_truncation = %(surrogate.first_pass_max_overweight_tru first_pass_freeze_max_weight_after = %(surrogate.first_pass_freeze_max_weight_after)s param_card = %(surrogate.param_card)s # for a .so surrogate, "" uses Cards/param_card.dat +[amplicol] +# the [amplicol] processes are integrated at leading colour; their events are +# re- or unweighted to color_accuracy in a second pass (lc: no second pass) +color_accuracy = %(amplicol.color_accuracy)s +mode = %(amplicol.mode)s +# unweighting at leading colour (unweight mode); [generation] sets the second pass +first_pass_max_overweight_truncation = %(amplicol.first_pass_max_overweight_truncation)s +first_pass_freeze_max_weight_after = %(amplicol.first_pass_freeze_max_weight_after)s + [multiparticles] $multiparticles diff --git a/madgraph/iolibs/template_files/mg7/run_interface.py b/madgraph/iolibs/template_files/mg7/run_interface.py index 4ba4a80ade..bd71bbd364 100644 --- a/madgraph/iolibs/template_files/mg7/run_interface.py +++ b/madgraph/iolibs/template_files/mg7/run_interface.py @@ -198,7 +198,7 @@ def get_model(self): tools this class drives (MadSpin, reweight, the auto widths and the 'update dependent' pass of their card questions) work here too. """ - from madgraph.iolibs.template_files.mg7.launch import load_process_model + from madgraph.iolibs.template_files.mg7.model_widths import load_process_model return load_process_model(self.me_dir) def do_set(self, line, log=True): diff --git a/madgraph/iolibs/template_files/mg7/subprocess_amplicol.py b/madgraph/iolibs/template_files/mg7/subprocess_amplicol.py new file mode 100644 index 0000000000..807edefd08 --- /dev/null +++ b/madgraph/iolibs/template_files/mg7/subprocess_amplicol.py @@ -0,0 +1,213 @@ +"""Subprocesses with AmpliCol matrix elements ('[amplicol]' processes). + +Each subprocess is one colour ordering (UMAMI channel) of an AmpliCol provider, +integrated on its own at leading colour with a colour-ordered phase-space +mapping. A second pass re- or unweights its events to full colour with the +ratio of the full-colour to the leading-colour matrix element summed over all +colour orderings, at the helicity of the event. +""" + +from __future__ import annotations + +import os +from typing import TYPE_CHECKING + +from madgraph.iolibs.template_files.mg7.common import lhe_helicities, ms +from madgraph.iolibs.template_files.mg7.observables import clean_pids +from madgraph.iolibs.template_files.mg7.subprocess_madmatrix import ( + Channel, + MadgraphSubprocess, + PhaseSpace, +) + +if TYPE_CHECKING: + from madgraph.iolibs.template_files.mg7.madgraph_process import MadgraphProcess + + +def is_amplicol(meta: dict) -> bool: + return meta.get("backend") == "amplicol" + + +def amplicol_parameters(meta: dict, process: MadgraphProcess) -> dict[str, float]: + """The values of the mutable model parameters in the param card.""" + parameters = {} + for name, (block, code) in meta["parameters"].items(): + try: + parameters[name] = float( + process.param_card.get_value(block.lower(), tuple(code)) + ) + except (KeyError, ValueError): + continue + return parameters + + +def chain_topology(incoming_masses: list[float], outgoing_masses: list[float]): + """A t-channel chain without propagators, for the LHE completion.""" + n_out = len(outgoing_masses) + vertices = [["i0", "p0", "o0"]] + vertices += [[f"p{i - 1}", f"p{i}", f"o{i}"] for i in range(1, n_out - 1)] + vertices.append(["i1", f"p{n_out - 2}", f"o{n_out - 1}"]) + propagators = [ms.Propagator() for _ in range(n_out - 1)] + return ms.Topology( + ms.Diagram(incoming_masses, outgoing_masses, propagators, vertices) + ) + + +def amplicol_subproc_args(meta: dict, process: MadgraphProcess) -> ms.SubprocArgs: + """LHE completion of an AmpliCol subprocess: its events have no diagram.""" + return ms.SubprocArgs( + topologies=[chain_topology( + [process.get_mass(pid) for pid in clean_pids(meta["incoming"])], + [process.get_mass(pid) for pid in clean_pids(meta["outgoing"])], + )], + permutations=[[]], + diagram_indices=[[]], + diagram_color_indices=[[]], + color_flows=meta["color_flows"], + pdg_color_types={ + int(key): value for key, value in meta["pdg_color_types"].items() + }, + helicities=lhe_helicities(meta, process.run_card), + pdg_ids=[flavor["options"] for flavor in meta["flavors"]], + ) + + +class AmplicolSubprocess(MadgraphSubprocess): + def __init__(self, process, meta, subproc_id, unmerged_meta=None): + super().__init__(process, meta, subproc_id, unmerged_meta) + self.check_flavor_blind_cuts() + + def check_flavor_blind_cuts(self): + """AmpliCol groups the flavours and identical-particle orderings that + share a matrix element into one channel, so its cuts, built for the + first flavour, must select the same particles for all of them.""" + options = [ + clean_pids(pdgs) for flavor in self.meta["flavors"] + for pdgs in flavor["options"] + ] + for item in self.process.cut_data: + for group in item.observable_kwargs["select_pids"]: + if len({tuple(pid in group for pid in pdgs) for pdgs in options}) > 1: + raise RuntimeError( + f"the cut {item.observable_kwargs['name']} does not select " + f"the same particles in all flavours of the [amplicol] " + f"subprocess {self.meta['incoming']} > {self.meta['outgoing']}" + ) + + def load_matrix_elements(self, all_api_paths, api_path_formats, member_metas): + run_args = self.process.run_card["run"] + if any(device != "cpu" for device in self.process.device_types): + raise RuntimeError("[amplicol] processes are only supported on CPUs") + if run_args["dummy_matrix_element"]: + raise RuntimeError("[amplicol] processes need their matrix elements") + if self.process.run_card["gridpack"]["save_gridpack"]: + raise RuntimeError("[amplicol] processes do not support gridpacks yet") + self.supports_alpha_s = self.meta.get("supports_alpha_s", True) + if not self.supports_alpha_s and not self.process.run_card["beam"]["fixed_ren_scale"]: + raise RuntimeError( + "the AmpliCol matrix elements of this model cannot take the running " + "alpha_s, as the model has no SMINPUTS 3 parameter: use a fixed " + "renormalization scale" + ) + parameters = amplicol_parameters(self.meta, self.process) + self.channel_index = self.meta["umami_channel"] + self.matrix_elements = [ + self.load(self.meta["me_path"], self.meta["artifact"], parameters) + ] + self.surrogates = [] + self.ratio_matrix_element = None + if self.process.run_card["amplicol"]["color_accuracy"] != "lc": + self.ratio_matrix_element = self.load( + self.meta["ratio_path"], self.meta["ratio_artifact"], parameters + ) + + def load(self, path: str, artifact: str, parameters: dict): + # loaded once for all channels; each load reads the whole pyamplicol artifact + apis = self.process.amplicol_apis + if path not in apis: + for context in self.process.contexts: + apis[path] = context.load_matrix_element( + path, os.path.abspath(artifact), parameters + ) + return apis[path] + + def build_multi_channel_data(self): + return None + + def build_multichannel_phasespace(self) -> PhaseSpace: + mapping = ms.PhaseSpaceMapping( + self.incoming_masses + self.outgoing_masses, + self.process.e_cm, + leptonic=self.process.leptonic, + invariant_power=self.process.run_card["phasespace"]["invariant_power"], + mode=ms.PhaseSpaceMapping.color_ordered, + cuts=self.cuts, + color_order=self.meta["phasespace_order"], + ) + prefix = f"subproc{self.subproc_id}.amplicol" + discrete_sym, discrete_flavor = self.build_discrete( + 1, len(self.meta["flavors"]), prefix + ) + channel = Channel( + phasespace_mapping=mapping, + adaptive_mapping=self.build_vegas(mapping, prefix), + discrete_sym=discrete_sym, + discrete_flavor=discrete_flavor, + channel_weight_indices=[0], + name="C", + active_flavors=[], + ) + return PhaseSpace(mode="multichannel", channels=[channel], symfact=[None]) + + def build_flat_phasespace(self) -> PhaseSpace: + return self.build_multichannel_phasespace() + + def simplify_phasespace(self, multi_phasespace, cross_sections) -> PhaseSpace: + return multi_phasespace + + def integrand_matrix_elements(self) -> list: + # without an alpha_s input, the model value of the fixed scale is used + inputs = [ + ms.MatrixElement.momenta_in, + *([ms.MatrixElement.alpha_s_in] if self.supports_alpha_s else []), + ms.MatrixElement.flavor_in, + ms.MatrixElement.channel_in, + ms.MatrixElement.random_helicity_in, + ] + outputs = [ + ms.MatrixElement.matrix_element_out, + ms.MatrixElement.color_index_out, + ms.MatrixElement.helicity_index_out, + ] + return [ms.MatrixElement( + self.matrix_elements[0].index(), self.particle_count, inputs, outputs, 1, + True, self.process.me_frame, self.process.incoming_count, + self.channel_index, + )] + + def has_weighting_steps(self) -> bool: + return self.ratio_matrix_element is not None + + def weighting_steps(self) -> list: + """The second pass to full colour.""" + if self.ratio_matrix_element is None: + return [] + ratio = ms.MatrixElement( + self.ratio_matrix_element.index(), + self.particle_count, + [ms.MatrixElement.momenta_in, ms.MatrixElement.flavor_in, + ms.MatrixElement.channel_in, ms.MatrixElement.helicity_in], + [ms.MatrixElement.matrix_element_ratio_out], + 1, + False, + self.process.me_frame, + self.process.incoming_count, + self.channel_index, + ) + gen_args = self.process.run_card["generation"] + return [ms.WeightingStep( + [ratio], + mode=self.process.run_card["amplicol"]["mode"], + max_overweight_truncation=gen_args["max_overweight_truncation"], + freeze_max_weight_after=gen_args["freeze_max_weight_after"], + )] diff --git a/madgraph/iolibs/template_files/mg7/subprocess_madmatrix.py b/madgraph/iolibs/template_files/mg7/subprocess_madmatrix.py new file mode 100644 index 0000000000..25ad94d7ea --- /dev/null +++ b/madgraph/iolibs/template_files/mg7/subprocess_madmatrix.py @@ -0,0 +1,1159 @@ +"""Subprocesses with madmatrix matrix elements and MG diagram based phase space.""" + +from __future__ import annotations + +import os +import re +from dataclasses import dataclass, field +from typing import Literal, NamedTuple, TYPE_CHECKING + +from madgraph.iolibs.template_files.mg7.common import logger, me_parameters, ms +from madgraph.iolibs.template_files.mg7.custom_function_loader import load_surrogate +from madgraph.iolibs.template_files.mg7.observables import build_observable, clean_pids, decay_products, observable_pids, pid_is_qcd + +if TYPE_CHECKING: + from madgraph.iolibs.template_files.mg7.madgraph_process import MadgraphProcess + +@dataclass +class Channel: + phasespace_mapping: ms.PhaseSpaceMapping + adaptive_mapping: ms.Flow | ms.VegasMapping + discrete_sym: ms.DiscreteSampler | ms.DiscreteFlow | None + discrete_flavor: ms.DiscreteSampler | ms.DiscreteFlow | None + channel_weight_indices: list[int] | None + name: str + active_flavors: list[int] + event_generator: ms.ChannelEventGenerator | None = None + + +@dataclass +class PhaseSpace: + mode: Literal["multichannel", "flat", "both"] + channels: list[Channel] + symfact: list[int | None] + first_chan_weight_remap: list[list[int]] = field(default_factory=list) + first_remapped_chan_count: int = 0 + second_chan_weight_remap: list[int] = field(default_factory=list) + second_remapped_chan_count: int = 0 + prop_chan_weights: ms.PropagatorChannelWeights | None = None + subchan_weights: ms.SubchannelWeights | None = None + cwnet: ms.ChannelWeightNetwork | None = None + + +class MultiChannelData(NamedTuple): + amp2_remaps: list[list[int]] + symfact: list[int | None] + topologies: list[list[ms.Topology]] + permutations: list[list[list[int]]] + channel_indices: list[list[int]] + channel_weight_indices: list[list[list[int]]] + diagram_indices: list[list[int]] + diagram_color_indices: list[list[list[int]]] + diagram_propagator_pdgs: list[list[list[int]]] + active_flavors: list[list[list[int]]] + qcd_s_channel_count: list[int] + + +def build_topologies( + incoming_masses: list[float], + outgoing_masses: list[float], + channel: dict, + process: MadgraphProcess +) -> list[ms.Topology]: + propagators = [] + for i, (pid, signed_pid) in enumerate(zip( + clean_pids(channel["propagators"]), channel["propagators"] + )): + mass = process.get_mass(pid) + width = process.get_width(pid) + bw_cutoff = process.run_card["phasespace"]["bw_cutoff"] + if i in channel["on_shell_propagators"]: + e_min = mass - bw_cutoff * width + e_max = mass + bw_cutoff * width + else: + e_min = 0 + e_max = 0 + # $-excluded propagator: the matrix element vanishes within bw_cutoff + # widths of the pole, so do not pile the channel's points up there + # (a flat density in that window instead of the Breit-Wigner peak). + # Channels written before the exporter recorded them have none. + flat_window = 0. + if i in channel.get("dollar_propagators", ()) and width > 0: + flat_window = float(bw_cutoff) + propagators.append(ms.Propagator( + mass=mass, + width=width, + integration_order=0, + e_min=e_min, + e_max=e_max, + pdg_id=signed_pid, + flat_window=flat_window, + )) + vertices = channel["vertices"] + diag = ms.Diagram( + incoming_masses, outgoing_masses, propagators, vertices + ) + return ms.Topology.topologies(diag) + + +def build_multi_channel_data( + meta: dict, process: MadgraphProcess, unmerged_meta: dict | None = None +) -> MultiChannelData: + incoming_masses = [ + process.get_mass(pid) for pid in clean_pids(meta["incoming"]) + ] + outgoing_masses = [ + process.get_mass(pid) for pid in clean_pids(meta["outgoing"]) + ] + + if unmerged_meta is None: + diagram_count = meta["diagram_count"] + amp2_remaps = [[-1] * diagram_count] + else: + amp2_remaps = [ + [-1] * unmerged_meta[subproc]["diagram_count"] + for subproc in meta["subprocesses"] + ] + symfact = [] + topologies = [] + permutations = [] + channel_indices = [] + channel_weight_indices = [] + diagram_indices = [] + diagram_color_indices = [] + diagram_propagator_pdgs = [] + active_flavors = [] + channel_index = 0 + qcd_s_channel_count = [] + + for channel in meta["channels"]: + if unmerged_meta is None: + topo_channel = channel + else: + topo_subproc = channel["subprocess"] + topo_channel_index = channel["channel"] + topo_channel = unmerged_meta[topo_subproc]["channels"][topo_channel_index] + chan_topologies = build_topologies( + incoming_masses, outgoing_masses, topo_channel, process + ) + topo_count = len(chan_topologies) + if topo_count == 0: + continue + + topo = chan_topologies[0] + s_chan_count = [0] * len(topo.decays) + non_qcd = [False] * len(topo.decays) + pdg_ids = [decay.pdg_id for decay in topo.decays] + for i, pid in zip(topo.outgoing_indices, meta["outgoing"]): + pdg_ids[i] = pid + for decay in reversed(topo.decays): + if len(decay.child_indices) == 0: + continue + if decay.index == 0 and topo.t_propagator_count > 0: + s_chan_count[0] = sum(s_chan_count[i] for i in decay.child_indices) + break + + is_qcd = pid_is_qcd(pdg_ids[decay.index]) and all( + pid_is_qcd(pdg_ids[i]) for i in decay.child_indices + ) + is_non_qcd = decay.on_shell or not is_qcd or any( + non_qcd[i] for i in decay.child_indices + ) + non_qcd[decay.index] = is_non_qcd + s_chan_count[decay.index] = ( + 0 if is_non_qcd else sum(s_chan_count[i] for i in decay.child_indices) + 1 + ) + qcd_s_channel_count.append(s_chan_count[0]) + + diagrams = channel["diagrams"] + chan_permutations = [d["permutation"] for d in diagrams] + if unmerged_meta is None: + amp2_remaps[0][diagrams[0]["diagram"]] = channel_index + else: + for amp2_remap, diag in zip(amp2_remaps, diagrams[0]["diagram"]): + if diag != -1: + amp2_remap[diag] = channel_index + + channel_index_first = channel_index + symfact_index_first = len(symfact) + channel_index += 1 + symfact.extend([None] * topo_count) + for d in diagrams[1:]: + if unmerged_meta is None: + amp2_remaps[0][d["diagram"]] = channel_index + else: + for amp2_remap, diag in zip(amp2_remaps, d["diagram"]): + if diag != -1: + amp2_remap[diag] = channel_index + channel_index += 1 + symfact.extend(range(symfact_index_first, symfact_index_first + topo_count)) + + topologies.append(chan_topologies) + permutations.append(chan_permutations) + channel_indices.append(list(range(channel_index_first, channel_index))) + channel_weight_indices.append([ + [ + symfact_index_first + topo_index + i * topo_count + for i in range(len(chan_permutations)) + ] + for topo_index in range(topo_count) + ]) + diagram_indices.append([d["diagram"] for d in diagrams]) + if unmerged_meta is None: + diagram_color_indices.append([d["active_colors"] for d in diagrams]) + diagram_propagator_pdgs.append( + [d["propagator_pdgs"] for d in diagrams] + ) + active_flavors.append([d["active_flavors"] for d in diagrams]) + + return MultiChannelData( + amp2_remaps, + symfact, + topologies, + permutations, + channel_indices, + channel_weight_indices, + diagram_indices, + diagram_color_indices, + diagram_propagator_pdgs, + active_flavors, + qcd_s_channel_count, + ) + + +class MadgraphSubprocess: + def __init__( + self, + process: MadgraphProcess, + meta: dict, + subproc_id: int, + unmerged_meta: dict | None = None + ): + self.process = process + self.meta = meta + self.subproc_id = subproc_id + self.multi_channel_data = None + + self.unmerged_meta = None + if unmerged_meta is None: + api_path_formats = [self.meta["me_path"]] + member_metas = [self.meta] + else: + api_path_formats = [] + member_metas = [] + for subproc in self.meta["subprocesses"]: + submeta = unmerged_meta[subproc] + api_path_formats.append(submeta["me_path"]) + member_metas.append(submeta) + if len(api_path_formats) == 1: + self.meta = submeta + else: + self.unmerged_meta = unmerged_meta + + # The libraries were all built up front by MadgraphProcess.compile_matrix_elements + all_api_paths = [ + [api_path_format.format(device=backend) for backend in self.process.backends] + for api_path_format in api_path_formats + ] + + self.incoming_masses = [ + self.process.get_mass(pid) for pid in clean_pids(self.meta["incoming"]) + ] + self.outgoing_masses = [ + self.process.get_mass(pid) for pid in clean_pids(self.meta["outgoing"]) + ] + self.particle_count = len(self.incoming_masses) + len(self.outgoing_masses) + # MadEvent's cut_decays = F: the decay products of on-shell propagators + # are left out of every cut selection (a pair cut then needs both + # members to be cut), while the histograms still see them + if self.process.run_card["phasespace"]["cut_decays"]: + uncut = set() + else: + uncut = decay_products(self.meta, self.unmerged_meta) + n_in = len(self.incoming_masses) + + def cut_pids(observable_kwargs): + pids = observable_pids(self.meta, observable_kwargs, self.unmerged_meta) + # pdg id 0 is in no group, so the particle is never selected + return [0 if i - n_in in uncut else pid for i, pid in enumerate(pids)] + + # custom cuts and scales see the PDG ids of the first flavor, like the + # built-in cuts, which select by flavor-independent groups + representative_pids = [self.flavor_pids()[0]] + custom_cut = self.process.custom_function( + "cut", self.particle_count, representative_pids + ) + if len(self.process.cut_data) > 0: + self.cuts = ms.Cuts([ + ms.CutItem( + observable=build_observable( + cut_pids(cut_item.observable_kwargs), + n_in, cut_item.observable_kwargs + ), + min=cut_item.min, + max=cut_item.max, + mode=cut_item.mode, + ) + for cut_item in self.process.cut_data + ], custom_cut) + elif custom_cut is not None: + self.cuts = ms.Cuts(self.particle_count, custom_cut) + else: + self.cuts = None + # the integration histograms are functions of the momenta; the weight + # distribution is a property of the final event sample, so it is only + # filled by MadgraphProcess.build_event_histograms at combine time + # They cost an observable evaluation per integration point, and the + # plots/HwU are drawn from the post-processing ones instead, so they + # are only filled on request ([run] weighted_histograms). + momentum_hists = [item for item in self.process.hist_data + if not item.from_weight] \ + if self.process.run_card["run"]["weighted_histograms"] else [] + self.histograms = ( + ms.ObservableHistograms([ + ms.HistItem( + observable=build_observable( + observable_pids( + self.meta, hist_item.observable_kwargs, self.unmerged_meta + ), + n_in, hist_item.observable_kwargs + ), + min=hist_item.min, + max=hist_item.max, + bin_count=hist_item.bin_count, + ) + for hist_item in momentum_hists + ]) + if momentum_hists + else None + ) + + self.scale = ms.EnergyScale( + particle_count=self.particle_count, + custom_scale=self.process.custom_function( + "scale", self.particle_count, representative_pids + ), + **self.process.scale_kwargs + ) + + self.load_matrix_elements(all_api_paths, api_path_formats, member_metas) + + def load_matrix_elements(self, all_api_paths, api_path_formats, member_metas): + """The matrix element and [surrogate] of every member subprocess.""" + if self.process.run_card["run"]["dummy_matrix_element"]: + self.matrix_elements = [None] * len(all_api_paths) + else: + self.matrix_elements = [] + for api_paths in all_api_paths: + for context, api_path in zip(self.process.contexts, api_paths): + mat = context.load_matrix_element( + api_path, self.process.param_card_path, + me_parameters(self.process.run_card) + ) + self.matrix_elements.append(mat) + self.surrogates = [] + if self.process.surrogate is not None: + for path_format, member_meta in zip(api_path_formats, member_metas): + self.surrogates.append(self.load_surrogate(path_format, member_meta)) + + def load_surrogate(self, me_path_format: str, member_meta: dict): + """The [surrogate] of one matrix element of the subprocess.""" + name = re.sub( + r"^libmadmatrix_(.*)_\{device\}\.so$", r"\1", os.path.basename(me_path_format) + ) + # PDG ids for each flavor index of the matrix element + flavors = member_meta["flavors"] + pdg_ids = [flavors[0]["options"][0]] * (max(f["index"] for f in flavors) + 1) + for flavor in flavors: + pdg_ids[flavor["index"]] = flavor["options"][0] + settings = self.process.surrogate + apis = load_surrogate( + ms, + self.process.contexts, + self.process.backends, + self.process.device_types, + settings["function"], + name, + pdg_ids, + settings["param_card"], + ) + if len({api.index() for api in apis}) != 1: + raise RuntimeError("surrogate loaded at different indices") + return apis[0] + + def has_weighting_steps(self) -> bool: + return bool(self.surrogates) + + def weighting_steps(self) -> list: + """The steps after the first unweighting: with a [surrogate], the + second pass with the full matrix element.""" + if not self.surrogates: + return [] + gen_args = self.process.run_card["generation"] + return [ms.WeightingStep( + self.step_matrix_elements, + mode=self.process.surrogate["mode"], + max_overweight_truncation=gen_args["max_overweight_truncation"], + freeze_max_weight_after=gen_args["freeze_max_weight_after"], + )] + + def flavor_pids(self) -> list[list[int]]: + """PDG ids of the external particles for every flavor of the + integrand, in the order of its flavor index.""" + pids = [] + for flav in self.meta["flavors"]: + if self.unmerged_meta is not None: + subproc_index = self.meta["subprocesses"][flav["subprocess"]] + flav = self.unmerged_meta[subproc_index]["flavors"][flav["flavor"]] + pids.append(flav["options"][0]) + return pids + + def build_multi_channel_data(self) -> MultiChannelData: + if self.multi_channel_data is not None: + return self.multi_channel_data + self.multi_channel_data = build_multi_channel_data( + self.meta, self.process, self.unmerged_meta + ) + return self.multi_channel_data + + def build_multichannel_phasespace(self) -> PhaseSpace: + mcdata = self.build_multi_channel_data() + channel_count = sum(len(topos) for topos in mcdata.topologies) + drop_threshold = self.process.run_card["phasespace"]["drop_qcd_s_channel"] + if drop_threshold >= 0 and channel_count > drop_threshold: + mcdata = self.drop_qcd_s_channels(mcdata) + mcdata = self.drop_empty_channels(mcdata) + + channels = [] + t_channel_mode = self.t_channel_mode( + self.process.run_card["phasespace"]["t_channel"] + ) + for channel_id, (chan_topologies, chan_permutations, chan_indices, active_flavors) in enumerate(zip( + mcdata.topologies, mcdata.permutations, mcdata.channel_weight_indices, + mcdata.active_flavors + )): + topo_count = len(chan_topologies) + for topo_index, (topo, indices) in enumerate(zip(chan_topologies, chan_indices)): + mapping = ms.PhaseSpaceMapping( + chan_topologies[0], + self.process.e_cm, + t_channel_mode=t_channel_mode, + cuts=self.cuts, + invariant_power=self.process.run_card["phasespace"]["invariant_power"], + permutations=chan_permutations, + leptonic=self.process.leptonic, + ) + prefix = f"subproc{self.subproc_id}.channel{channel_id}" + if topo_count > 1: + prefix += f".subchan{topo_index}" + discrete_sym, discrete_flavor = self.build_discrete( + len(chan_permutations), len(self.meta["flavors"]), prefix + ) + channels.append(Channel( + phasespace_mapping = mapping, + adaptive_mapping = self.build_vegas(mapping, prefix), + discrete_sym = discrete_sym, + discrete_flavor = discrete_flavor, + channel_weight_indices = indices, + name = f"{channel_id}", + active_flavors = active_flavors, + )) + + remapped_chan_count = sum( + len(indices) for indices in mcdata.channel_indices + ) + if self.process.run_card["phasespace"]["sde_strategy"] == "denominators": + prop_chan_weights = ms.PropagatorChannelWeights( + [topo[0] for topo in mcdata.topologies], mcdata.permutations, + mcdata.channel_indices + ) + chan_weight_remap = [] + else: + prop_chan_weights = None + chan_weight_remap = [ + [ + len(mcdata.symfact) if remap == -1 else remap + for remap in amp2_remap + ] + for amp2_remap in mcdata.amp2_remaps + ] + + if any(len(topos) > 1 for topos in mcdata.topologies): + subchan_weights = ms.SubchannelWeights( + mcdata.topologies, mcdata.permutations, mcdata.channel_indices + ) + else: + subchan_weights = None + + return PhaseSpace( + mode="multichannel", + channels=channels, + first_chan_weight_remap=chan_weight_remap, + first_remapped_chan_count=remapped_chan_count, + symfact=mcdata.symfact, + prop_chan_weights=prop_chan_weights, + subchan_weights=subchan_weights, + ) + + def build_flat_phasespace(self) -> PhaseSpace: + mapping = ms.PhaseSpaceMapping( + self.incoming_masses + self.outgoing_masses, + self.process.e_cm, + mode=self.t_channel_mode(self.process.run_card["phasespace"]["flat_mode"]), + cuts=self.cuts, + leptonic=self.process.leptonic, + ) + prefix = f"subproc{self.subproc_id}.flat" + discrete_sym, discrete_flavor = self.build_discrete( + 1, len(self.meta["flavors"]), prefix + ) + channel = Channel( + phasespace_mapping = mapping, + adaptive_mapping = self.build_vegas(mapping, prefix), + discrete_sym = discrete_sym, + discrete_flavor = discrete_flavor, + channel_weight_indices = [0], + name = "F", + active_flavors = [], + ) + if self.unmerged_meta is None: + remap = [list(range(self.meta["diagram_count"]))] + else: + remap = [ + list(range(self.unmerged_meta[subproc]["diagram_count"])) + for subproc in self.meta["subprocesses"] + ] + return PhaseSpace( + mode="flat", + channels=[channel], + first_chan_weight_remap=remap, + first_remapped_chan_count=1, + symfact=[None], + ) + + def simplify_phasespace( + self, + multi_phasespace: PhaseSpace, + cross_sections: list[float] + ) -> PhaseSpace | None: + assert multi_phasespace.mode == "multichannel" + + threshold = 1 - self.process.run_card["phasespace"]["combine_channel_threshold"] + kept_channels = [] + tot_cs = sum(cross_sections) + cum_cs = 0. + seen_active_flavors = set() + #seen_resonances = set() + for index, (cs, chan) in sorted( + enumerate(zip(cross_sections, multi_phasespace.channels)), + key=lambda pair: pair[1][0], + reverse=True + ): + cum_cs += cs + has_unseen_flavors = False + has_unseen_resonances = False + for flavs in chan.active_flavors: + for flav in flavs: + if flav not in seen_active_flavors: + has_unseen_flavors = True + seen_active_flavors.add(flav) + #for resonance in chan.resonances: + # if resonance not in seen_resonances: + # has_unseen_resonances = True + # seen_resonances.add(flav) + if has_unseen_flavors or has_unseen_resonances or not tot_cs \ + or cum_cs / tot_cs < threshold: + kept_channels.append(index) + if len(kept_channels) >= len(cross_sections) - 1: + return multi_phasespace + + channels = [] + channel_map = {} + symfact = [] + for old_chan_index in kept_channels: + channel = multi_phasespace.channels[old_chan_index] + perm_count = max(1, channel.phasespace_mapping.channel_count()) + channel_index = len(symfact) + symfact.append(None) + symfact.extend([channel_index] * (perm_count - 1)) + channel_map.update({ + old_index: new_index + for new_index, old_index in enumerate( + channel.channel_weight_indices, start=channel_index + ) + }) + channels.append(Channel( + phasespace_mapping = channel.phasespace_mapping, + adaptive_mapping = channel.adaptive_mapping, + discrete_sym = channel.discrete_sym, + discrete_flavor = channel.discrete_flavor, + channel_weight_indices = list(range( + channel_index, channel_index + perm_count + )), + name = channel.name, + active_flavors = channel.active_flavors, + event_generator = channel.event_generator, + )) + + flat_phasespace = self.build_flat_phasespace() + flat_channel = flat_phasespace.channels[0] + channels.append(Channel( + phasespace_mapping = flat_channel.phasespace_mapping, + adaptive_mapping = flat_channel.adaptive_mapping, + discrete_sym = flat_channel.discrete_sym, + discrete_flavor = flat_channel.discrete_flavor, + channel_weight_indices = [len(symfact)], + name = flat_channel.name, + active_flavors = flat_channel.active_flavors, + )) + flat_index = len(symfact) + symfact.append(None) + channel_map[len(multi_phasespace.symfact)] = len(symfact) + if multi_phasespace.subchan_weights is None and len(multi_phasespace.first_chan_weight_remap) > 0: + first_chan_weight_remap = [ + [ + channel_map.get(remap, flat_index) + for remap in cw_remap + ] + for cw_remap in multi_phasespace.first_chan_weight_remap + ] + first_remapped_chan_count = len(symfact) + second_chan_weight_remap = [] + second_remapped_chan_count = 0 + else: + first_chan_weight_remap = multi_phasespace.first_chan_weight_remap + first_remapped_chan_count = multi_phasespace.first_remapped_chan_count + chan_count = multi_phasespace.first_remapped_chan_count if multi_phasespace.subchan_weights is None else multi_phasespace.subchan_weights.channel_count() + second_chan_weight_remap = [ + channel_map.get(i, flat_index) + for i in range(chan_count) + ] + second_remapped_chan_count = len(symfact) + + return PhaseSpace( + mode="both", + channels=channels, + first_chan_weight_remap=first_chan_weight_remap, + first_remapped_chan_count=first_remapped_chan_count, + second_chan_weight_remap=second_chan_weight_remap, + second_remapped_chan_count=second_remapped_chan_count, + symfact=symfact, + prop_chan_weights=multi_phasespace.prop_chan_weights, + subchan_weights=multi_phasespace.subchan_weights, + ) + + def drop_qcd_s_channels(self, mcdata: MultiChannelData) -> MultiChannelData: + """Drop channels with non-resonant QCD s-channel propagators, to reduce the + channel count for processes with many diagrams. Channels are grouped by the + number of QCD s-channels, allowing for more if necessary to map out all flavor + indices. Channel weights belonging to a dropped channel are left unmapped.""" + groups_by_s_count = {} + for index, count in enumerate(mcdata.qcd_s_channel_count): + groups_by_s_count.setdefault(count, []).append(index) + + covered_flavors = set() + kept_groups = set() + for s_count in sorted(groups_by_s_count): + s_count_groups = groups_by_s_count[s_count] + if s_count == 0: + selected = s_count_groups + else: + selected = [ + index + for index in s_count_groups + if any( + flav not in covered_flavors + for flavs in mcdata.active_flavors[index] + for flav in flavs + ) + ] + kept_groups.update(selected) + for index in selected: + for flavs in mcdata.active_flavors[index]: + covered_flavors.update(flavs) + + return self.keep_channel_groups(mcdata, sorted(kept_groups)) + + def drop_empty_channels(self, mcdata: MultiChannelData) -> MultiChannelData: + """Drop the channels whose phase space the cuts exclude entirely: a pair + mass cut whose floor lies at or above the top of an on-shell window of + the channel (ms.PhaseSpaceMapping.empty). Every point such a channel + generates fails the cuts, so it could only ever fail to find one. + + If that leaves a flavour with no channel at all, its cross section is + zero -- typically a cut that excludes the window of a decay chain every + diagram shares -- and the run cannot sample it: that is reported as an + error naming the cause, rather than as the survey finding no point + passing the cuts.""" + if self.cuts is None: + return mcdata + t_channel_mode = self.t_channel_mode( + self.process.run_card["phasespace"]["t_channel"] + ) + empty = [ + ms.PhaseSpaceMapping( + chan_topologies[0], + self.process.e_cm, + t_channel_mode=t_channel_mode, + cuts=self.cuts, + invariant_power=self.process.run_card["phasespace"]["invariant_power"], + permutations=chan_permutations, + leptonic=self.process.leptonic, + ).empty() + for chan_topologies, chan_permutations in zip( + mcdata.topologies, mcdata.permutations + ) + ] + if not any(empty): + return mcdata + kept_groups = [index for index, is_empty in enumerate(empty) if not is_empty] + covered_flavors = { + flav + for index in kept_groups + for flavs in mcdata.active_flavors[index] + for flav in flavs + } + uncovered = sorted({ + flav + for index, is_empty in enumerate(empty) if is_empty + for flavs in mcdata.active_flavors[index] + for flav in flavs + } - covered_flavors) + if uncovered: + raise ValueError( + f"subprocess {self.meta['incoming']} > {self.meta['outgoing']}: " + "the pair mass cuts exclude the whole on-shell window " + f"(bw_cutoff = {self.process.run_card['phasespace']['bw_cutoff']}) " + f"of every channel of flavour(s) {uncovered}, so their cross " + "section is zero; loosen the cut or remove that decay" + ) + logger.info( + "subprocess %d: %d channel(s) dropped, the cuts exclude the " + "on-shell window they sample", self.subproc_id, len(empty) - len(kept_groups) + ) + return self.keep_channel_groups(mcdata, kept_groups) + + def keep_channel_groups( + self, mcdata: MultiChannelData, kept_groups: list[int] + ) -> MultiChannelData: + """Rebuild the multi-channel data with only the channel groups (indices + into mcdata.topologies) in kept_groups. Channel weights belonging to a + dropped channel are left unmapped.""" + if len(kept_groups) == len(mcdata.topologies): + return mcdata + + amp2_remaps = [[-1] * len(remap) for remap in mcdata.amp2_remaps] + symfact = [] + topologies = [] + permutations = [] + channel_indices = [] + channel_weight_indices = [] + diagram_indices = [] + diagram_color_indices = [] + diagram_propagator_pdgs = [] + active_flavors = [] + qcd_s_channel_count = [] + channel_index = 0 + + for group in kept_groups: + chan_topologies = mcdata.topologies[group] + chan_permutations = mcdata.permutations[group] + chan_diagram_indices = mcdata.diagram_indices[group] + topo_count = len(chan_topologies) + + channel_index_first = channel_index + symfact_index_first = len(symfact) + for i, diag in enumerate(chan_diagram_indices): + if self.unmerged_meta is None: + amp2_remaps[0][diag] = channel_index + else: + for amp2_remap, d in zip(amp2_remaps, diag): + if d != -1: + amp2_remap[d] = channel_index + if i == 0: + symfact.extend([None] * topo_count) + else: + symfact.extend(range(symfact_index_first, symfact_index_first + topo_count)) + channel_index += 1 + + topologies.append(chan_topologies) + permutations.append(chan_permutations) + channel_indices.append(list(range(channel_index_first, channel_index))) + channel_weight_indices.append([ + [ + symfact_index_first + topo_index + i * topo_count + for i in range(len(chan_permutations)) + ] + for topo_index in range(topo_count) + ]) + diagram_indices.append(chan_diagram_indices) + if mcdata.diagram_color_indices: + diagram_color_indices.append(mcdata.diagram_color_indices[group]) + if mcdata.diagram_propagator_pdgs: + diagram_propagator_pdgs.append( + mcdata.diagram_propagator_pdgs[group] + ) + active_flavors.append(mcdata.active_flavors[group]) + qcd_s_channel_count.append(mcdata.qcd_s_channel_count[group]) + + return MultiChannelData( + amp2_remaps, + symfact, + topologies, + permutations, + channel_indices, + channel_weight_indices, + diagram_indices, + diagram_color_indices, + diagram_propagator_pdgs, + active_flavors, + qcd_s_channel_count, + ) + + def set_madnis_auto_settings(self, rsd: float): + madnis_args = self.process.run_card["madnis"] + n_out = len(self.meta["outgoing"]) + n_events = self.process.run_card["generation"]["events"] + is_gridpack = self.process.run_card["gridpack"]["save_gridpack"] + train_batches = madnis_args["train_batches"] + hidden_dim = min(max(int((7 * rsd) / 32) * 32 + 64, 64), 256) + flow_layers = 3 if rsd < 32 else 4 + lr = min(max((10000 - train_batches) / 8000 * 7e-4 + 3e-4, 3e-4), 1e-3) + if n_out <= 2: + enable = False + elif n_out == 3: + enable = rsd > 10. or n_events > 1000000 or is_gridpack + else: + enable = True + fixed_cwnet_fraction = max(0.33, 1.0 - 10000. / train_batches) + batch_size_per_channel = min(max(int((7 * rsd) / 32) * 32 + 64, 128), 512) + + self.madnis_settings = { + "enable": enable, + "flow_layers": flow_layers, + "flow_hidden_dim": hidden_dim, + "discrete_hidden_dim": hidden_dim, + "cwnet_hidden_dim": hidden_dim, + "train_batches": train_batches, + "lr": lr, + "fixed_cwnet_fraction": fixed_cwnet_fraction, + "batch_size_per_channel": batch_size_per_channel, + } + for key, value in self.madnis_settings.items(): + run_card_value = madnis_args[key] + if run_card_value != "auto": + self.madnis_settings[key] = run_card_value + + def build_madnis(self, phasespace: PhaseSpace) -> PhaseSpace: + madnis_args = self.process.run_card["madnis"] + # Shared across all networks below: initialize_globals() derives an + # independent, non-colliding stream per tensor from this one base seed. + seed = self.process.run_seed + channels = [] + for channel_id, channel in enumerate(phasespace.channels): + prefix = f"subproc{self.subproc_id}.channel{channel_id}" + cond_dim = 0 + + flow_dim = channel.phasespace_mapping.random_dim() + flow = ms.Flow( + input_dim=flow_dim, + condition_dim=cond_dim, + prefix=prefix, + bin_count=madnis_args["flow_spline_bins"], + subnet_hidden_dim=self.madnis_settings["flow_hidden_dim"], + subnet_layers=self.madnis_settings["flow_layers"], + subnet_activation=self.activation(madnis_args["flow_activation"]), + invert_spline=madnis_args["flow_invert_spline"], + ) + if channel.adaptive_mapping is None: + flow.initialize_globals(self.process.contexts[0], seed) + else: + flow.initialize_from_vegas( + self.process.contexts[0], channel.adaptive_mapping.grid_name(), + seed + ) + cond_dim += flow_dim + + # discrete_sym runs after the adaptive map, so it can condition on its latent. + discrete_sym = channel.discrete_sym + if discrete_sym is not None: + perm_count = channel.phasespace_mapping.channel_count() + discrete_sym = ms.DiscreteFlow( + option_counts=[perm_count], + prefix=f"{prefix}.discrete_flow_sym", + dims_with_prior=[], + condition_dim=cond_dim, + subnet_hidden_dim=self.madnis_settings["discrete_hidden_dim"], + subnet_layers=madnis_args["discrete_layers"], + subnet_activation=self.activation(madnis_args["discrete_activation"]), + ) + discrete_sym.initialize_globals(self.process.contexts[0], seed) + cond_dim += perm_count + + discrete_flavor = channel.discrete_flavor + if discrete_flavor is not None: + discrete_flavor = ms.DiscreteFlow( + option_counts=[len(self.meta["flavors"])], + prefix=f"{prefix}.discrete_flow_flavor", + dims_with_prior=[0], + condition_dim=cond_dim, + subnet_hidden_dim=self.madnis_settings["discrete_hidden_dim"], + subnet_layers=madnis_args["discrete_layers"], + subnet_activation=self.activation(madnis_args["discrete_activation"]), + ) + discrete_flavor.initialize_globals(self.process.contexts[0], seed) + + channels.append(Channel( + phasespace_mapping = channel.phasespace_mapping, + adaptive_mapping = flow, + discrete_sym = discrete_sym, + discrete_flavor = discrete_flavor, + channel_weight_indices = channel.channel_weight_indices, + name = channel.name, + active_flavors = channel.active_flavors, + )) + + return PhaseSpace( + mode="both", + channels=channels, + first_chan_weight_remap=phasespace.first_chan_weight_remap, + first_remapped_chan_count=phasespace.first_remapped_chan_count, + second_chan_weight_remap=phasespace.second_chan_weight_remap, + second_remapped_chan_count=phasespace.second_remapped_chan_count, + symfact=phasespace.symfact, + cwnet=self.build_cwnet(len(phasespace.symfact)), + prop_chan_weights=phasespace.prop_chan_weights, + subchan_weights=phasespace.subchan_weights, + ) + + def build_vegas(self, mapping: ms.PhaseSpaceMapping, prefix: str) -> ms.VegasMapping: + if not self.process.run_card["vegas"]["enable"]: + return None + + vegas = ms.VegasMapping( + mapping.random_dim(), + self.process.run_card["vegas"]["bins"], + prefix, + ) + for context in self.process.contexts: + vegas.initialize_globals(context) + return vegas + + def build_discrete( + self, permutation_count: int, flavor_count: int, prefix: str + ) -> tuple[ms.DiscreteSampler | None, ms.DiscreteSampler | None]: + is_adaptive = self.process.run_card["phasespace"]["adaptive_symmetry_sampling"] + if is_adaptive and permutation_count > 1: + discrete_sym = ms.DiscreteSampler( + [permutation_count], f"{prefix}.discrete_sym" + ) + for context in self.process.contexts: + discrete_sym.initialize_globals(context) + else: + discrete_sym = None + + if flavor_count > 1: + # the flavor probabilities are conditioned on the PDF prior; a + # leptonic process (or a decay) has none, and the Integrand then + # passes no condition: declaring one anyway fails to build + # ("keys and values must have the same size"), e.g. z > q q~ with + # both u- and d-type rows + discrete_flavor = ms.DiscreteSampler( + [flavor_count], f"{prefix}.discrete_flavor", + [] if self.process.leptonic else [0] + ) + for context in self.process.contexts: + discrete_flavor.initialize_globals(context) + else: + discrete_flavor = None + + return discrete_sym, discrete_flavor + + def build_cwnet(self, channel_count: int) -> ms.ChannelWeightNetwork: + #if channel_count == 1: + # return None + madnis_args = self.process.run_card["madnis"] + cwnet = ms.ChannelWeightNetwork( + channel_count=channel_count, + particle_count=self.particle_count, + hidden_dim=self.madnis_settings["cwnet_hidden_dim"], + layers=madnis_args["cwnet_layers"], + activation=self.activation(madnis_args["cwnet_activation"]), + prefix=f"subproc{self.subproc_id}.cwnet", + ) + cwnet.initialize_globals( + self.process.contexts[0], self.process.run_seed + ) + return cwnet + + def t_channel_mode(self, name: str) -> ms.PhaseSpaceMapping.TChannelMode: + modes = { + "propagator": ms.PhaseSpaceMapping.propagator, + "rambo": ms.PhaseSpaceMapping.rambo, + "chili": ms.PhaseSpaceMapping.chili, + } + if name in modes: + return modes[name] + else: + raise ValueError(f"Invalid t-channel mode '{name}'") + + def activation(self, name: str) -> ms.MLP.Activation: + activations = { + "relu": ms.MLP.relu, + "leaky_relu": ms.MLP.leaky_relu, + "elu": ms.MLP.elu, + "gelu": ms.MLP.gelu, + "sigmoid": ms.MLP.sigmoid, + "softplus": ms.MLP.softplus, + } + if name in activations: + return activations[name] + else: + raise ValueError(f"Invalid activation function '{name}'") + + def integrand_matrix_elements(self) -> list: + """The matrix element of every member subprocess in the integrand. Sets + the step_matrix_elements of the second pass of a [surrogate].""" + matrix_elements = [] + self.step_matrix_elements = [] + for me_index, matrix_element in enumerate(self.matrix_elements): + if self.surrogates: + # the surrogate in the integrand, the full matrix element in + # the second pass, with the random numbers passed in + mat = ms.MatrixElement( + self.surrogates[me_index].index(), + self.particle_count, + [ms.MatrixElement.momenta_in, ms.MatrixElement.alpha_s_in, + ms.MatrixElement.flavor_in], + [ms.MatrixElement.matrix_element_out], + 1, + False, + self.process.me_frame, + self.process.incoming_count, + ) + self.step_matrix_elements.append(ms.MatrixElement( + matrix_element, + [ms.MatrixElement.momenta_in, ms.MatrixElement.alpha_s_in, + ms.MatrixElement.flavor_in, ms.MatrixElement.random_color_in, + ms.MatrixElement.random_helicity_in, + ms.MatrixElement.random_diagram_in], + [ms.MatrixElement.matrix_element_out, + ms.MatrixElement.color_index_out, + ms.MatrixElement.helicity_index_out, + ms.MatrixElement.diagram_index_out], + False, + self.process.me_frame, + self.process.incoming_count, + )) + elif matrix_element: + mat = ms.MatrixElement( + matrix_element, + ms.Integrand.matrix_element_inputs, + ms.Integrand.matrix_element_outputs, + True, + self.process.me_frame, + self.process.incoming_count, + ) + else: + #TODO: not working in merged mode + mat = ms.MatrixElement( + 0xBADCAFE, + self.particle_count, + ms.Integrand.matrix_element_inputs, + ms.Integrand.matrix_element_outputs, + self.meta["diagram_count"], + True, + self.process.me_frame, + self.process.incoming_count, + ) + matrix_elements.append(mat) + return matrix_elements + + def build_integrands( + self, + phasespace: PhaseSpace, + madnis_training: bool = False, + drop_cuts_and_rescale: bool = False + ) -> list[ms.Integrand]: + flavors = [] + flavor_remap = [] + flavor_factors = [] + flavor_mirror = [] + flavor_diff_xs_indices = [] + flavor_subproc_indices = [] + flavor_per_subproc_remap = [] + + for flav in self.meta["flavors"]: + if self.unmerged_meta is not None: + diff_xs_index = flav["subprocess"] + subproc_index = self.meta["subprocesses"][diff_xs_index] + ps_flavor = flav["flavor"] + flavor_diff_xs_indices.append(diff_xs_index) + flavor_subproc_indices.append(subproc_index) + flavor_per_subproc_remap.append(ps_flavor) + flav = self.unmerged_meta[subproc_index]["flavors"][ps_flavor] + flavors.append(flav["options"][0]) + flavor_remap.append(flav["index"]) + flavor_factors.append(flav.get("factor", len(flav["options"]))) + flavor_mirror.append(flav["mirror"]) + + cross_sections = [] + for mat in self.integrand_matrix_elements(): + pdf_grid = None if self.process.leptonic else self.process.pdf_grid + pdf_arg = None if self.process.leptonic else ms.CachedPdf() + cross_sections.append( + ms.DifferentialCrossSection( + matrix_element=mat, + # 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(), + pid_options=[], + pdf1=pdf_arg, + pdf2=pdf_arg, + input_momentum_fraction=True, + decay=self.process.is_decay, + return_matrix_element=self.has_weighting_steps(), + ) + ) + partial_weights = self.process.systematics_enabled() + madnis_args = self.process.run_card["madnis"] + custom_bias = self.process.custom_function("bias", self.particle_count, flavors) + integrands = [] + for channel in phasespace.channels: + integrands.append(ms.Integrand( + channel.phasespace_mapping, + cross_sections, + channel.adaptive_mapping, + channel.discrete_sym, + channel.discrete_flavor, + flavors, + pdf_grid, + self.process.running_coupling, + self.scale, + phasespace.prop_chan_weights, + phasespace.subchan_weights, + phasespace.cwnet, + phasespace.first_chan_weight_remap, + phasespace.first_remapped_chan_count, + phasespace.second_chan_weight_remap, + phasespace.second_remapped_chan_count, + madnis_training, + drop_cuts_and_rescale, + partial_weights, + channel.channel_weight_indices, + channel.active_flavors, + flavor_remap, + flavor_factors, + flavor_mirror, + flavor_diff_xs_indices, + flavor_subproc_indices, + flavor_per_subproc_remap, + madnis_args["compressed_channel_weight_count"], + custom_bias, + )) + #print(integrands[1].function()) + #for i in integrands: print(i.function()) + return integrands diff --git a/madgraph/various/banner.py b/madgraph/various/banner.py index c904f8691d..bcc6c6fb68 100755 --- a/madgraph/various/banner.py +++ b/madgraph/various/banner.py @@ -6917,6 +6917,17 @@ def default_setup(self): comment="freeze_max_weight_after of the unweighting with the surrogate") self.add_toml_param('surrogate', 'param_card', "", comment="param card for a .so surrogate; \"\" uses Cards/param_card.dat") + # ---------------------------- [amplicol] ---------------------- + # the leading-colour [amplicol] subprocesses are re- or unweighted to + # the colour accuracy in a second pass + self.add_toml_param('amplicol', 'color_accuracy', "full", allowed=["full", "lc"], + comment="colour accuracy of the [amplicol] events, lc skips the second pass") + self.add_toml_param('amplicol', 'mode', "reweight", allowed=["unweight", "reweight"], + comment="unweight or reweight the events to the colour accuracy") + self.add_toml_param('amplicol', 'first_pass_max_overweight_truncation', 0.01, + comment="max_overweight_truncation of the leading-colour unweighting") + self.add_toml_param('amplicol', 'first_pass_freeze_max_weight_after', 10000, + comment="freeze_max_weight_after of the leading-colour unweighting") # ----------------------------- [madnis] ----------------------- self.add_toml_param('madnis', 'enable', False, allowed=["auto", True, False], auto=True, diff --git a/madmatrix/output.py b/madmatrix/output.py index 6e9cc06a88..7b299e5576 100644 --- a/madmatrix/output.py +++ b/madmatrix/output.py @@ -216,6 +216,12 @@ def get_makefile_replace_dict(self, model): # AV - overload the default version: create CMake directory, do not create lib directory def copy_template(self, model): + if not getattr(self, 'has_madmatrix', True): + # only the cards and the launcher: no src/, backend/ or makefiles + self.dirs_to_create = ['bin', 'lib', 'Cards', 'SubProcesses'] + self.from_template = {key: self.from_template[key] for key in ('.', 'Cards')} + self.template_src_make = self.template_Sub_make = None + return super().copy_template(model) super().copy_template(model) # Copy Arithmetics headers for the double-word expansion (FPTYPE=e) arithmetics_src = pjoin(self.madmatrix_templates, 'Arithmetics') @@ -293,6 +299,8 @@ def generate_subprocess_directory(self, matrix_element, cpp_helas_call_writer, p def convert_model(self, model, wanted_lorentz=[], wanted_couplings=[], **opts): # **opts: an option meant for one exporter (npwave, for the dual HELAS # libraries of P-wave bound states) is passed by keyword to all of them + if not getattr(self, 'has_madmatrix', True): + return if hasattr(model , 'cudacpp_wanted_ordered_couplings'): wanted_couplings = model.cudacpp_wanted_ordered_couplings del model.cudacpp_wanted_ordered_couplings diff --git a/madspace/include/madspace/driver/lhe_output.hpp b/madspace/include/madspace/driver/lhe_output.hpp index fde1385ea2..d580554edb 100644 --- a/madspace/include/madspace/driver/lhe_output.hpp +++ b/madspace/include/madspace/driver/lhe_output.hpp @@ -162,7 +162,8 @@ class LHECompleter { struct SubprocArgs { /// Process id shared by every event of this subprocess. int process_id; - /// Decay topology of each diagram channel. + /// Decay topology of each diagram channel. A subprocess without diagrams + /// has one topology and no permutations and diagram indices. std::vector topologies; /// Final-state particle permutation of each channel. nested_vector3 permutations; @@ -197,7 +198,8 @@ class LHECompleter { * @param event The event to complete, modified in place. * @param subprocess_index Index into the `subproc_args` passed to the * constructor. - * @param diagram_index Sampled diagram. + * @param diagram_index Sampled diagram, -1 for a subprocess without + * diagrams. * @param color_index Sampled color configuration. * @param flavor_index Sampled flavor assignment. * @param helicity_index Sampled helicity configuration. diff --git a/madspace/include/madspace/phasespace/integrand.hpp b/madspace/include/madspace/phasespace/integrand.hpp index 3da4e6e9c8..1783bdca30 100644 --- a/madspace/include/madspace/phasespace/integrand.hpp +++ b/madspace/include/madspace/phasespace/integrand.hpp @@ -94,7 +94,9 @@ class Integrand : public FunctionGenerator { }; /** - * @param mapping The phase-space mapping. + * @param mapping The phase-space mapping. Its + * discrete choices, if any, are + * sampled uniformly. * @param diff_xs One differential cross section * per active flavour. * @param adaptive_map Optional continuous adaptive diff --git a/madspace/include/madspace/phasespace/matrix_element.hpp b/madspace/include/madspace/phasespace/matrix_element.hpp index 8d8a5a72f5..e30b442b14 100644 --- a/madspace/include/madspace/phasespace/matrix_element.hpp +++ b/madspace/include/madspace/phasespace/matrix_element.hpp @@ -98,6 +98,8 @@ class MatrixElement : public FunctionGenerator { * not commute, so going there straight from * the lab frame would leave a Wigner rotation * behind and rotate the polarisation axes. + * @param channel If non-negative, `channel_in` is passed this + * fixed value instead of being an input. */ MatrixElement( std::size_t matrix_element_index, @@ -107,7 +109,8 @@ class MatrixElement : public FunctionGenerator { std::size_t diagram_count = 1, bool sample_random_inputs = false, const std::vector& me_frame = {}, - std::size_t incoming_count = 2 + std::size_t incoming_count = 2, + me_int_t channel = -1 ); /** * Take @p matrix_element_index, @p particle_count and @p diagram_count from @@ -116,6 +119,9 @@ class MatrixElement : public FunctionGenerator { * @param inputs Quantities to pass in. * @param outputs Quantities to return. * @param sample_random_inputs See the other constructor. + * @param me_frame See the other constructor. + * @param incoming_count See the other constructor. + * @param channel See the other constructor. */ MatrixElement( const MatrixElementApi& matrix_element_api, @@ -123,7 +129,8 @@ class MatrixElement : public FunctionGenerator { const std::vector& outputs = {matrix_element_out}, bool sample_random_inputs = false, const std::vector& me_frame = {}, - std::size_t incoming_count = 2 + std::size_t incoming_count = 2, + me_int_t channel = -1 ) : MatrixElement( matrix_element_api.index(), @@ -133,7 +140,8 @@ class MatrixElement : public FunctionGenerator { matrix_element_api.diagram_count(), sample_random_inputs, me_frame, - incoming_count + incoming_count, + channel ) {}; /// Index of the matrix element in the plugin. std::size_t matrix_element_index() const { return _matrix_element_index; } @@ -147,8 +155,10 @@ class MatrixElement : public FunctionGenerator { const std::vector& outputs() const { return _outputs; } /// Whether the `random_*` inputs are drawn internally. bool sample_random_inputs() const { return _sample_random_inputs; } + /// The fixed channel passed as `channel_in`, or -1 if it is an input. + me_int_t channel() const { return _channel; } /// The inputs that must be supplied by the caller (excludes the ones drawn - /// internally when @p sample_random_inputs is set). + /// internally when @p sample_random_inputs is set and a fixed channel). std::vector external_inputs() const; /// Per-particle 0/1 selector built from me_frame, empty if no boost is /// applied. @@ -169,6 +179,7 @@ class MatrixElement : public FunctionGenerator { std::vector _inputs; std::vector _outputs; bool _sample_random_inputs; + me_int_t _channel; std::vector _frame_mask; std::vector _reference_mask; }; diff --git a/madspace/src/driver/lhe_output.cpp b/madspace/src/driver/lhe_output.cpp index 1e586ac138..245715e552 100644 --- a/madspace/src/driver/lhe_output.cpp +++ b/madspace/src/driver/lhe_output.cpp @@ -529,7 +529,10 @@ void LHECompleter::complete_event_data( if (event.particles.size() != subproc_data.particle_count) { throw std::runtime_error("Invalid particle number for subprocess"); } - if (diagram_index < 0 || diagram_index >= subproc_data.diagram_count) { + // a subprocess without diagrams has no propagators to insert + bool has_diagrams = subproc_data.diagram_count > 0; + if (has_diagrams ? diagram_index < 0 || diagram_index >= subproc_data.diagram_count + : diagram_index != -1) { throw std::runtime_error("Diagram index out of range"); } if (color_index < 0 || color_index >= subproc_data.color_count) { @@ -583,6 +586,9 @@ void LHECompleter::complete_event_data( ++particle_index; } + if (!has_diagrams) { + return; + } auto find_propagators = _propagator_index_and_count.find( cantor_pairing(subprocess_index, diagram_index, color_index) ); @@ -619,8 +625,7 @@ void LHECompleter::complete_event_data( // negative m_min would instead turn it into a large positive floor and // reject nearly everything, so a wider window would write *fewer* // resonances than a narrow one. - double m_min = - std::max(0., propagator.mass - _bw_cutoff * propagator.width); + double m_min = std::max(0., propagator.mass - _bw_cutoff * propagator.width); double m_max = propagator.mass + _bw_cutoff * propagator.width; if (m2 > m_min * m_min && m2 < m_max * m_max) { auto [color, anti_color] = prop_color; diff --git a/madspace/src/phasespace/integrand.cpp b/madspace/src/phasespace/integrand.cpp index 5e9893e8ef..868d5cfe5a 100644 --- a/madspace/src/phasespace/integrand.cpp +++ b/madspace/src/phasespace/integrand.cpp @@ -504,8 +504,22 @@ NamedVector Integrand::build_channel_part( chan_index_in_group = fb.full({static_cast(0), batch_size_val}); } - // Apply phase space mapping - auto mapping_result = _mapping.build_forward(fb, {latent}, mapping_conditions); + // Apply phase space mapping, its discrete choices are sampled uniformly + ValueVec mapping_inputs{latent}; + if (std::size_t discrete_dim = _mapping.discrete_dim(); discrete_dim > 0) { + ValueVec choices; + Value discrete_random = + fb.random(batch_size_val, static_cast(discrete_dim)); + for (auto& random : fb.unstack(discrete_random)) { + auto [choice, choice_det] = + fb.sample_discrete(random, static_cast(2)); + choices.push_back(choice); + weights_before_cuts.push_back(choice_det); + } + mapping_inputs.push_back(fb.stack(choices)); + } + auto mapping_result = + _mapping.build_forward(fb, mapping_inputs, mapping_conditions); weights_before_cuts.push_back(mapping_result["det"]); Value momenta = mapping_result["momenta"]; Value x0 = mapping_result["x1"]; diff --git a/madspace/src/phasespace/matrix_element.cpp b/madspace/src/phasespace/matrix_element.cpp index c0ef852d4c..5f869136e8 100644 --- a/madspace/src/phasespace/matrix_element.cpp +++ b/madspace/src/phasespace/matrix_element.cpp @@ -40,7 +40,8 @@ MatrixElement::MatrixElement( std::size_t diagram_count, bool sample_random_inputs, const std::vector& me_frame, - std::size_t incoming_count + std::size_t incoming_count, + me_int_t channel ) : FunctionGenerator( "MatrixElement", @@ -81,7 +82,9 @@ MatrixElement::MatrixElement( arg_types.push_back("diagram", batch_int); break; case channel_in: - arg_types.push_back("channel", batch_int); + if (channel < 0) { + arg_types.push_back("channel", batch_int); + } break; case particle_count_in: arg_types.push_back("particle_count", batch_int); @@ -159,6 +162,7 @@ MatrixElement::MatrixElement( _particle_count(particle_count), _diagram_count(diagram_count), _sample_random_inputs(sample_random_inputs), + _channel(channel), _frame_mask( me_frame.empty() ? std::vector() : particle_mask(particle_count, me_frame, "me_frame") @@ -274,6 +278,8 @@ NamedVector MatrixElement::build_function_impl( input == random_diagram_in)) { matrix_args.push_back(random.at(random_index)); ++random_index; + } else if (input == channel_in && _channel >= 0) { + matrix_args.push_back(fb.full({_channel, fb.batch_size(args.values())})); } else if (input == momenta_in && !_frame_mask.empty()) { // Evaluate the matrix element in the frame the run card asked for // (me_frame). Only the matrix element sees the boosted momenta: @@ -346,6 +352,9 @@ NamedVector MatrixElement::build_function_impl( std::vector MatrixElement::external_inputs() const { std::vector ret; for (auto input : _inputs) { + if (input == channel_in && _channel >= 0) { + continue; + } if (!_sample_random_inputs || (input != random_color_in && input != random_helicity_in && input != random_diagram_in)) { diff --git a/madspace/src/phasespace/weighting_step.cpp b/madspace/src/phasespace/weighting_step.cpp index c0e75f7075..8982b1bf24 100644 --- a/madspace/src/phasespace/weighting_step.cpp +++ b/madspace/src/phasespace/weighting_step.cpp @@ -107,7 +107,16 @@ WeightingStepFunction::WeightingStepFunction( ); } } - if (matrix_element.external_inputs().size() != matrix_element.inputs().size()) { + bool has_random_inputs = std::any_of( + matrix_element.inputs().begin(), + matrix_element.inputs().end(), + [](auto input) { + return input == MatrixElement::random_color_in || + input == MatrixElement::random_helicity_in || + input == MatrixElement::random_diagram_in; + } + ); + if (has_random_inputs && matrix_element.sample_random_inputs()) { throw std::invalid_argument( "the random inputs of a weighting step matrix element are passed in " "and must not be sampled internally" diff --git a/madspace/src/python/madspace.cpp b/madspace/src/python/madspace.cpp index 42ed9d1c83..8f39f4f12b 100644 --- a/madspace/src/python/madspace.cpp +++ b/madspace/src/python/madspace.cpp @@ -1564,7 +1564,8 @@ PYBIND11_MODULE(_madspace_py, m) { std::size_t, bool, const std::vector&, - std::size_t>(), + std::size_t, + me_int_t>(), py::arg("matrix_element_index"), py::arg("particle_count"), // C++ defaults to {momenta_in} / {matrix_element_out}; kept required @@ -1576,6 +1577,7 @@ PYBIND11_MODULE(_madspace_py, m) { py::arg("sample_random_inputs") = false, py::arg("me_frame") = std::vector{}, py::arg("incoming_count") = 2, + py::arg("channel") = -1, pydoc::doc("MatrixElement::MatrixElement") ) .def( @@ -1585,13 +1587,15 @@ PYBIND11_MODULE(_madspace_py, m) { const std::vector&, bool, const std::vector&, - std::size_t>(), + std::size_t, + me_int_t>(), py::arg("matrix_element_api"), py::arg("inputs"), py::arg("outputs"), py::arg("sample_random_inputs") = false, py::arg("me_frame") = std::vector{}, py::arg("incoming_count") = 2, + py::arg("channel") = -1, pydoc::doc("MatrixElement::MatrixElement#2") ) .def( @@ -1611,6 +1615,7 @@ PYBIND11_MODULE(_madspace_py, m) { ) .def("inputs", &MatrixElement::inputs, pydoc::doc("MatrixElement::inputs")) .def("outputs", &MatrixElement::outputs, pydoc::doc("MatrixElement::outputs")) + .def("channel", &MatrixElement::channel, pydoc::doc("MatrixElement::channel")) .def( "external_inputs", &MatrixElement::external_inputs, diff --git a/madspace/tests/test_integrand_me_outputs.py b/madspace/tests/test_integrand_me_outputs.py index 81b745bdf4..a5a9fc5c53 100644 --- a/madspace/tests/test_integrand_me_outputs.py +++ b/madspace/tests/test_integrand_me_outputs.py @@ -21,11 +21,11 @@ import madspace as ms -def python_matrix_element(context, function, outputs=("matrix_element",)): +def python_matrix_element(context, function, outputs=("matrix_element",), channel=-1): + """With a channel, the function also gets the fixed channel index.""" + inputs = ["momenta"] if channel < 0 else ["momenta", "channel_index"] api = context.add_matrix_element( - ms.PythonFunction( - function, ["momenta"], list(outputs), particle_count=len(PIDS) - ) + ms.PythonFunction(function, inputs, list(outputs), particle_count=len(PIDS)) ) output_enums = { "matrix_element": ms.MatrixElement.matrix_element_out, @@ -37,8 +37,10 @@ def python_matrix_element(context, function, outputs=("matrix_element",)): return ms.MatrixElement( api.index(), len(PIDS), - [ms.MatrixElement.momenta_in], + [ms.MatrixElement.momenta_in] + + ([] if channel < 0 else [ms.MatrixElement.channel_in]), [output_enums[output] for output in outputs], + channel=channel, ) @@ -176,3 +178,51 @@ def test_diagram_channel_weights_need_amp2(tmp_path): first_chan_weight_remap=[[0, 1]], first_remapped_chan_count=2, ) + + +@pytest.mark.parametrize("color_order", [[0, 2, 3, 4, 1], [0, 1, 2, 3, 4]]) +def test_color_ordered_discrete_choices(tmp_path, color_order): + """The integrand samples the discrete two-solution choices of a color-ordered + mapping, so it reproduces the phase-space volume of RAMBO.""" + context, alphas_file = make_context(str(tmp_path), 2) + matrix_element = python_matrix_element( + context, lambda momenta: np.ones(len(momenta)) + ) + cuts = ms.Cuts( + [ + ms.CutItem(ms.Observable(PIDS, ms.Observable.obs_pt, [[i]]), min=20.0) + for i in (2, 3, 4) + ] + ) + results = [] + for mode, order in [ + (ms.PhaseSpaceMapping.color_ordered, color_order), + (ms.PhaseSpaceMapping.rambo, None), + ]: + mapping = ms.PhaseSpaceMapping( + [0.0] * len(PIDS), + CM_ENERGY, + leptonic=True, + mode=mode, + cuts=cuts, + color_order=order, + ) + if order is not None: + assert mapping.discrete_dim() > 0 + diff_xs = ms.DifferentialCrossSection( + matrix_element=matrix_element, + cm_energy=CM_ENERGY, + running_coupling=None, + energy_scale=ms.CachedScale(), + ) + integrand = ms.Integrand( + mapping, + [diff_xs], + running_coupling=ms.RunningCoupling(ms.AlphaSGrid(alphas_file)), + energy_scale=ms.EnergyScale(len(PIDS), 91.188), + channel_indices=[0], + ) + results.append(generate(str(tmp_path), context, integrand, str(mode))[0]) + status, rambo_status = results + error = np.hypot(status.error, rambo_status.error) + assert abs(status.mean - rambo_status.mean) < 5 * error diff --git a/madspace/tests/test_lhe.py b/madspace/tests/test_lhe.py index c1c14fb07f..9df98bfbe8 100644 --- a/madspace/tests/test_lhe.py +++ b/madspace/tests/test_lhe.py @@ -461,3 +461,47 @@ def test_particle_momentum_kwargs_match_attributes(): for name in ("p_x", "p_y", "p_z"): with pytest.raises(TypeError): ms.LHEParticle(**{name: 1.0}) + + +def chain_topology(masses): + """t-channel chain topology without propagator information.""" + n_out = len(masses) - 2 + vertices = [["i0", "p0", "o0"]] + vertices += [[f"p{i - 1}", f"p{i}", f"o{i}"] for i in range(1, n_out - 1)] + vertices.append(["i1", f"p{n_out - 2}", f"o{n_out - 1}"]) + diagram = ms.Diagram( + masses[:2], masses[2:], [ms.Propagator() for _ in range(n_out - 1)], vertices + ) + return ms.Topology(diagram) + + +def test_subprocess_without_diagrams(): + """Without diagrams, events carry the diagram index -1 and get no + intermediate particles.""" + pdg_ids = [21, 21, 21, 21, 21] + color_flows = [ + [[502, 501], [503, 502], [503, 504], [505, 501], [505, 504]], + [[501, 502], [502, 503], [504, 503], [501, 505], [504, 505]], + ] + subproc_args = ms.SubprocArgs( + process_id=3, + topologies=[chain_topology([0.0] * 5)], + permutations=[[]], + diagram_indices=[[]], + diagram_color_indices=[[]], + color_flows=color_flows, + pdg_color_types={21: 8}, + helicities=[[1, 1, -1, -1, 1], [-1, -1, 1, 1, -1]], + pdg_ids=[[pdg_ids]], + ) + lhe_completer = ms.LHECompleter([subproc_args], bw_cutoff=BW_CUTOFF) + event = build_event(np.ones((5, 4)) * [100.0, 0.0, 0.0, 0.0]) + lhe_completer.complete_event_data(event, 0, -1, 1, 0, 1, ms.MixMaxRandom(1)) + assert event.process_id == 3 + assert [p.pdg_id for p in event.particles] == pdg_ids + assert [(p.color, p.anti_color) for p in event.particles] == [ + tuple(flow) for flow in color_flows[1] + ] + assert [p.spin for p in event.particles] == [-1, -1, 1, 1, -1] + with pytest.raises(RuntimeError, match="Diagram index out of range"): + lhe_completer.complete_event_data(event, 0, 0, 0, 0, 0, ms.MixMaxRandom(1)) diff --git a/madspace/tests/test_python_function.py b/madspace/tests/test_python_function.py index 4f290637d5..924094938a 100644 --- a/madspace/tests/test_python_function.py +++ b/madspace/tests/test_python_function.py @@ -91,6 +91,33 @@ def matrix_element(p): assert (np.asarray(helicity) == 3).all() +def test_fixed_channel(context, momenta): + def matrix_element(p, flavor, channel): + return p[:, 2, 0] * (channel + 10 * flavor) + + api = context.add_matrix_element( + ms.PythonFunction( + matrix_element, + ["momenta", "flavor_index", "channel_index"], + ["matrix_element"], + particle_count=4, + ) + ) + inputs = [ + ms.MatrixElement.momenta_in, + ms.MatrixElement.flavor_in, + ms.MatrixElement.channel_in, + ] + generator = ms.MatrixElement( + api.index(), 4, inputs, [ms.MatrixElement.matrix_element_out], channel=3 + ) + assert generator.channel() == 3 + assert generator.external_inputs() == inputs[:2] + flavor = np.full(len(momenta), 2, dtype=np.int32) + me = ms.FunctionRuntime(generator.function(), context)(momenta, flavor) + assert np.allclose(me, momenta[:, 2, 0] * 23) + + @pytest.mark.parametrize( "result, message", [ diff --git a/madspace/tests/test_weighting_steps.py b/madspace/tests/test_weighting_steps.py index fef4cbcf8b..1967cc74f8 100644 --- a/madspace/tests/test_weighting_steps.py +++ b/madspace/tests/test_weighting_steps.py @@ -76,6 +76,7 @@ def generate( step_outputs=("matrix_element",), override_outputs=(), integrand_outputs=("matrix_element",), + step_channel=-1, ): """Survey and generate, returns the status, the combined events and the channels. step_me and step_outputs can be lists, one entry per step.""" @@ -86,7 +87,7 @@ def generate( step_me = [step_me] * len(step_modes) step_outputs = [step_outputs] * len(step_modes) step_mes = [ - python_matrix_element(context, me, outputs) + python_matrix_element(context, me, outputs, step_channel) for me, outputs in zip(step_me, step_outputs) ] config = make_config() @@ -256,6 +257,32 @@ def test_ratio_step(tmp_path, step_modes): assert status_me.mean == status_ratio.mean +def test_ratio_with_fixed_channel(tmp_path): + """A step matrix element evaluated at a fixed channel gets it as an input, + which is not an input of the step.""" + + def ratio(momenta, channel): + return true_over_surrogate(momenta) * (channel == 3) + + results = [ + generate( + str(tmp_path / name), + surrogate_me, + ["reweight"], + step_me=step_me, + step_outputs=("matrix_element_ratio",), + step_channel=channel, + ) + for name, step_me, channel in [ + ("plain", true_over_surrogate, -1), + ("channel", ratio, 3), + ] + ] + (status_plain, events_plain, _), (status_channel, events_channel, _) = results + assert np.array_equal(events_plain, events_channel) + assert status_plain.mean == status_channel.mean + + def test_ratio_then_matrix_element(tmp_path, reference): """After a ratio step, the matrix element it implies is the reference of the next step, so the ratios still multiply to |M_true|^2 / |M_surrogate|^2.""" diff --git a/tests/acceptance_tests/test_mg7_amplicol.py b/tests/acceptance_tests/test_mg7_amplicol.py new file mode 100644 index 0000000000..802aa2ca72 --- /dev/null +++ b/tests/acceptance_tests/test_mg7_amplicol.py @@ -0,0 +1,166 @@ +################################################################################ +# +# Copyright (c) 2026 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 +# +################################################################################ +"""mg7 runs of '[amplicol]' processes against the same processes with madmatrix. + +The AmpliCol subprocesses are integrated at leading colour and re- or +unweighted to full colour, so their cross sections must agree with the +madmatrix ones within the Monte Carlo errors, alone and mixed with madmatrix +processes in one output. Needs madspace, pyamplicol and LHAPDF; the tests skip +otherwise. Knobs: MG7_AMPLICOL_EVENTS (default 10000) and MG7_AMPLICOL_NSIGMA +(default 4). + + ./tests/test_manager.py test_mg7_amplicol -pA -t0 +""" + +import glob +import gzip +import importlib.util +import json +import os +import re +import shutil +import subprocess +import sys +import tempfile +import unittest + +import madgraph.interface.master_interface as MGCmd +import madgraph.various.misc as misc + +pjoin = os.path.join + +_EVENTS = int(os.environ.get('MG7_AMPLICOL_EVENTS', 10000)) +_NSIGMA = float(os.environ.get('MG7_AMPLICOL_NSIGMA', 4)) +_PDF = 'NNPDF23_lo_as_0130_qed' + + +def _edit_run_card(path, **amplicol): + card = open(path).read() + card = card.replace('fixed_ren_scale = false', 'fixed_ren_scale = true') + card = card.replace('fixed_fact_scale = false', 'fixed_fact_scale = true') + card = re.sub(r'(?m)^pdf = ".*"$', 'pdf = "%s"' % _PDF, card) + card = re.sub(r'(?m)^events = \d+', 'events = %d' % _EVENTS, card) + card = re.sub(r'(?m)^seed = -1', 'seed = 7', card) + card = re.sub(r'(?m)^make_plots = true', 'make_plots = false', card) + card = re.sub(r'(\[systematics\]\nenable = )true', r'\1false', card) + card = re.sub(r'(\[madnis\]\n(?:#.*\n)*enable = )"auto"', r'\1false', card) + head, section = card.split('\n[amplicol]\n') + for key, value in amplicol.items(): + section, count = re.subn(r'(?m)^%s = .*$' % key, + '%s = "%s"' % (key, value), section, count=1) + assert count == 1, key + card = head + '\n[amplicol]\n' + section + open(path, 'w').write(card) + + +def _color_balanced(lhe_path): + """Whether every colour line of every event starts and ends once.""" + with gzip.open(lhe_path, 'rt') as stream: + text = stream.read() + for event in re.findall(r'\n(.*?)', text, re.S): + lines = [line.split() for line in event.strip().splitlines()[1:] + if line.strip() and not line.startswith('<')] + tags = {} + for particle in lines: + status, color, anticolor = int(particle[1]), int(particle[4]), int(particle[5]) + if status == -1: + color, anticolor = anticolor, color + for tag, sign in ((color, 1), (anticolor, -1)): + if tag: + tags[tag] = tags.get(tag, 0) + sign + if any(tags.values()): + return False + return True + + +class MG7AmplicolTest(unittest.TestCase): + + def setUp(self): + # not imported: Symbolica would hold its lock against the pyamplicol worker + missing = [name for name in ('madspace', 'pyamplicol') + if importlib.util.find_spec(name) is None] + if missing: + self.skipTest('%s unavailable' % ' and '.join(missing)) + if not misc.resolve_lhapdf().find_set(_PDF): + self.skipTest('%s LHAPDF data not found' % _PDF) + self.path = tempfile.mkdtemp(prefix='mg7_amplicol_') + + def tearDown(self): + shutil.rmtree(self.path, ignore_errors=True) + + def run_mg7(self, name, processes, **amplicol): + """Output and run the processes, returns the cross section, its error + and the run directory.""" + run_dir = pjoin(self.path, name) + mg = MGCmd.MasterCmd() + mg.no_notification() + mg.exec_cmd('set automatic_html_opening False --no_save') + for index, process in enumerate(processes): + mg.exec_cmd('%s %s' % ('generate' if index == 0 else 'add process', + process)) + mg.exec_cmd('output mg7 %s' % run_dir) + _edit_run_card(pjoin(run_dir, 'Cards', 'run_card.toml'), **amplicol) + log = pjoin(run_dir, 'generate_events.log') + with open(log, 'w') as stream: + status = subprocess.call( + [sys.executable, pjoin(run_dir, 'bin', 'generate_events'), '-f'], + cwd=run_dir, stdout=stream, stderr=subprocess.STDOUT) + self.assertEqual(status, 0, open(log).read()[-3000:]) + info = sorted(glob.glob(pjoin(run_dir, 'Events', '*', 'info.json')))[-1] + with open(info) as stream: + result = json.load(stream)['process'] + return float(result['mean']), float(result['error']), run_dir + + def assert_compatible(self, first, second, factor=1.): + (mean1, error1), (mean2, error2) = first, second + error = (error1**2 + (factor * error2)**2)**0.5 + self.assertLess(abs(mean1 - factor * mean2), _NSIGMA * error, + '%g +- %g vs %g +- %g' % (mean1, error1, factor * mean2, + factor * error2)) + + def test_gg_ggg(self): + """Full colour in both modes, and leading colour, which differs by + (N^2 - 1) / N^2 for five gluons.""" + madmatrix = self.run_mg7('mm', ['g g > g g g'])[:2] + reweight = self.run_mg7('rw', ['g g > g g g [amplicol]']) + self.assert_compatible(reweight[:2], madmatrix) + events = glob.glob(pjoin(reweight[2], 'Events', '*', 'events.lhe.gz')) + if events: + self.assertTrue(_color_balanced(events[-1])) + unweight = self.run_mg7('uw', ['g g > g g g [amplicol]'], mode='unweight') + self.assert_compatible(unweight[:2], madmatrix) + leading = self.run_mg7('lc', ['g g > g g g [amplicol]'], + color_accuracy='lc') + self.assert_compatible(madmatrix, leading[:2], 8. / 9.) + + def test_pp_jjj(self): + """Quark lines and the expansion of the multiparticle labels.""" + madmatrix = self.run_mg7('mm', ['p p > j j j'])[:2] + amplicol = self.run_mg7('amp', ['p p > j j j [amplicol]'])[:2] + self.assert_compatible(amplicol, madmatrix) + + def test_gg_ttxg(self): + """Massive quarks in the colour-ordered phase space.""" + madmatrix = self.run_mg7('mm', ['g g > t t~ g'])[:2] + amplicol = self.run_mg7('amp', ['g g > t t~ g [amplicol]'])[:2] + self.assert_compatible(amplicol, madmatrix) + + def test_mixed(self): + """[amplicol] and madmatrix processes in one output.""" + jj = self.run_mg7('jj', ['p p > j j'])[:2] + ggg = self.run_mg7('ggg', ['g g > g g g'])[:2] + mixed = self.run_mg7('mixed', ['p p > j j', 'g g > g g g [amplicol]'])[:2] + self.assert_compatible( + mixed, (jj[0] + ggg[0], (jj[1]**2 + ggg[1]**2)**0.5)) diff --git a/tests/unit_tests/interface/test_amplicol_syntax.py b/tests/unit_tests/interface/test_amplicol_syntax.py new file mode 100644 index 0000000000..854b7bb3ff --- /dev/null +++ b/tests/unit_tests/interface/test_amplicol_syntax.py @@ -0,0 +1,153 @@ +################################################################################ +# +# Copyright (c) 2026 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 +# +################################################################################ +"""Syntax of the processes computed by AmpliCol + + generate p p > j j j [amplicol] + +The process is kept as a definition for the mg7 output, which generates its +matrix elements with pyamplicol; no MG diagrams are generated for it. +""" + +import os +import sys + +root_path = os.path.split(os.path.dirname(os.path.realpath(__file__)))[0] +sys.path.append(os.path.join(root_path, os.path.pardir, os.path.pardir)) + +import tests.unit_tests as unittest + +import madgraph.interface.madgraph_interface as mg_interface +import madgraph.interface.master_interface as MGCmd + +from madgraph import InvalidCmd + +split = mg_interface.split_backend_line + + +class BackendLineSplitTest(unittest.TestCase): + """split_backend_line, without any model.""" + + def test_split(self): + self.assertEqual(split('p p > j j j [amplicol]'), ('p p > j j j', 'amplicol')) + self.assertEqual(split('p p > j j j [ AmpliCol ] @2'), + ('p p > j j j @2', 'amplicol')) + self.assertEqual(split('p p > j j j'), ('p p > j j j', '')) + self.assertEqual(split('p p > t t~ [QCD]'), ('p p > t t~ [QCD]', '')) + + def test_other_brackets_are_refused(self): + self.assertRaises(InvalidCmd, split, 'p p > j j [amplicol] [QCD]') + + def test_process_type(self): + process_type = MGCmd.MasterCmd.extract_process_type + self.assertEqual(process_type('p p > j j j [amplicol]'), ('tree', None, [])) + self.assertEqual(process_type('p p > t t~ [QCD]'), ('NLO', 'all', ['QCD'])) + + +class AmplicolProcessTest(unittest.TestCase): + """generate/add process with [amplicol], with the sm model.""" + + def setUp(self): + self.cmd = MGCmd.MasterCmd() + self.cmd.exec_cmd('import model sm', printcmd=False, precmd=True) + + def test_no_diagrams_are_generated(self): + self.cmd.exec_cmd('generate g g > g g g [amplicol]', printcmd=False, + precmd=True) + self.cmd.exec_cmd('add process p p > e+ e-', printcmd=False, precmd=True) + self.assertEqual(len(self.cmd._curr_amplicol_defs), 1) + self.assertEqual( + self.cmd._curr_amplicol_defs[0].nice_string(prefix=False).split('@')[0].strip(), + 'g g > g g g') + # the madmatrix process keeps its own process number + self.assertTrue(self.cmd._curr_amps) + self.assertEqual( + {amp.get('process').get('id') for amp in self.cmd._curr_amps}, {2}) + self.cmd.exec_cmd('generate p p > e+ e-', printcmd=False, precmd=True) + self.assertEqual(len(self.cmd._curr_amplicol_defs), 0) + + def test_unsupported_options_are_refused(self): + for line in ['p p > j j QED=0 [amplicol]', + 'p p > z > e+ e- [amplicol]', + 'p p > j j / b [amplicol]', + 'z > e+ e- [amplicol]']: + self.assertRaises(InvalidCmd, self.cmd.do_generate, line) + + +class AmplicolExportTest(unittest.TestCase): + """The process definitions as pyamplicol sees them, without pyamplicol.""" + + def setUp(self): + self.cmd = MGCmd.MasterCmd() + self.cmd.exec_cmd('import model sm', printcmd=False, precmd=True) + self.cmd.exec_cmd('define q = u d', printcmd=False, precmd=True) + self.cmd.exec_cmd('define qx = u~ d~', printcmd=False, precmd=True) + + def export(self, line): + import madgraph.iolibs.export_amplicol as export_amplicol + self.cmd.exec_cmd('generate %s [amplicol]' % line, printcmd=False, + precmd=True) + procdef = self.cmd._curr_amplicol_defs[0] + model = self.cmd._curr_model + return (export_amplicol.process_name(procdef, model, + self.cmd._multiparticles), + export_amplicol.process_expression(procdef, model), + export_amplicol.representative_processes([procdef], model)) + + def test_flavour_groups_are_expanded(self): + """MG replaces the light quarks by a flavour group, which pyamplicol + does not know: the original flavours are passed on.""" + name, expression, processes = self.export('u u~ > e+ e- g') + self.assertEqual(name, 'uux_epemg') + self.assertEqual(expression, ('u u~ > e+ e- g', {})) + self.assertEqual(len(processes), 1) + name, expression, processes = self.export('q qx > e+ e- g') + self.assertEqual(name, 'qqx_epemg') + self.assertEqual(expression, ('mp0 mp1 > e+ e- g', + {'mp0': ['u', 'd'], 'mp1': ['u~', 'd~']})) + self.assertEqual(len(processes), 4) + + def test_multiparticles(self): + name, (expression, multiparticles), processes = self.export('p p > j j') + self.assertEqual(name, 'pp_jj') + self.assertEqual(expression, 'mp0 mp0 > mp0 mp0') + self.assertEqual(sorted(multiparticles['mp0']), + sorted(['g', 'u', 'd', 's', 'c', 'u~', 'd~', 's~', 'c~'])) + self.assertEqual(len(processes), 81) + + def test_coupling_orders(self): + """The orders MG selects by default, as the largest non-QCD orders.""" + import madgraph.iolibs.export_amplicol as export_amplicol + for line, orders in [('p p > j j j', {'QED': 0}), + ('p p > e+ e- j', {'QED': 2}), + ('p p > w+ j j', {'QED': 1})]: + self.cmd.exec_cmd('generate %s [amplicol]' % line, printcmd=False, + precmd=True) + self.assertEqual(export_amplicol.coupling_orders( + self.cmd._curr_amplicol_defs[0], self.cmd._curr_model), orders) + + def test_mirrored_beams(self): + """pyamplicol keeps one process per unordered initial state; the other + beam order is integrated as the mirror where the definition allows it.""" + import madgraph.iolibs.export_amplicol as export_amplicol + mirrored = export_amplicol.mirrored + partons = {21, 1, 2, -1, -2} + exported = {((21, 2), (2, 21, 21)), ((2, 2), (2, 2, 21)), + ((2, -2), (-2, 2, 21)), ((-2, 2), (-2, 2, 21))} + self.assertTrue(mirrored([21, 2, 21, 21, 2], 2, [partons, partons], exported)) + self.assertFalse(mirrored([2, 2, 2, 2, 21], 2, [partons, partons], exported)) + # both beam orders are exported + self.assertFalse(mirrored([2, -2, 2, -2, 21], 2, [partons, partons], exported)) + # g u > u g g only + self.assertFalse(mirrored([21, 2, 21, 21, 2], 2, [{21}, {2}], exported)) diff --git a/tests/unit_tests/interface/test_mg7_launch.py b/tests/unit_tests/interface/test_mg7_launch.py index 8f7fd39788..12f851e099 100644 --- a/tests/unit_tests/interface/test_mg7_launch.py +++ b/tests/unit_tests/interface/test_mg7_launch.py @@ -846,14 +846,15 @@ def test_setup_logging_defers_to_an_existing_handler(self): lg = logging.getLogger('madgraph') handler = logging.NullHandler() lg.addHandler(handler) - saved = self.launch._TOOL_LOGGING_READY - self.launch._TOOL_LOGGING_READY = False + from madgraph.iolibs.template_files.mg7 import postprocess + saved = postprocess._TOOL_LOGGING_READY + postprocess._TOOL_LOGGING_READY = False try: before = list(lg.handlers) self.launch._setup_logging() self.assertEqual(list(lg.handlers), before) finally: - self.launch._TOOL_LOGGING_READY = saved + postprocess._TOOL_LOGGING_READY = saved lg.removeHandler(handler)