npsim/scripts/sanitize_hepmc3.py

0001 #!/usr/bin/env python3
0002
0003 import sys
0004 import argparse
0005
0006 class InvalidHepmc3Error(Exception):
0007     def __init__(this, *msg):
0008         print('InvalidHepmc3Error:', *msg, file=sys.stderr)
0009
0010 def warn(*msg):
0011     print('WARNING:', *msg, file=sys.stderr)
0012
0013
0014 class EventHeader:
0015     def __init__(this, line):
0016         if not line[0] == 'E':
0017             raise InvalidHepmc3Error('Expected event record, but found:', line)
0018         this.raw = line
0019         this.vertex = this._get_vertex_if_present(line)
0020         this.n_vertices = int(line.split(' ')[2])
0021         this.n_particles = int(line.split(' ')[3])
0022         this.vert_cnt = 0
0023         this.part_cnt = 0
0024     def get_record(this):
0025         if this.vert_cnt > this.n_vertices:
0026             warn('Too many vertices for event:', this.raw, "--> skipping event")
0027             return None
0028         if not this.part_cnt == this.n_particles:
0029             warn('Invalid particle count for event:', this.raw, "--> skipping event")
0030             return None
0031         if this.vertex and not '@' in this.raw:
0032             return '{} @{}'.format(this.raw[:-1], this.vertex)
0033         return this.raw
0034     def process_vertex(this, line):
0035         this.vert_cnt += 1
0036         if this.vert_cnt > this.n_vertices:
0037             warn('Too many vertices for event:', this.raw)
0038             return
0039         if not this.vertex:
0040             if '@' in line:
0041                 this.vertex = this._get_vertex_if_present(line)
0042     def process_particle(this, line):
0043         this.part_cnt += 1
0044         if this.part_cnt > this.n_particles:
0045             warn('Too many particles for event:', this.raw)
0046             return
0047     def _get_vertex_if_present(this, line):
0048         if not '@' in line:
0049             return None
0050         return line.split('@')[1]
0051
0052 def flush_buffer(header, buffer):
0053     if header:
0054         event_record = header.get_record()
0055         if not event_record:
0056             warn('Skipped invalid event', header.raw)
0057             return
0058         sys.stdout.write(header.get_record())
0059     for line in buffer:
0060         sys.stdout.write(line)
0061
0062 if __name__ == '__main__':
0063     parser = argparse.ArgumentParser(
0064             description='Read HepMC3 input from stdin and sanitize the '
0065                         'output (e.g. add vertex info to the event header '
0066                         'if the event has a displaced starting vertex). '
0067                         'Output is written to stdout.')
0068     args = parser.parse_args()
0069
0070     header = None
0071     buffer = []
0072     have_first_line = False
0073     have_second_line = False
0074     end_reached = False
0075     ## read line-by-line, fill the event buffer, and then
0076     ## write the sanitized output to stdout
0077     for line in sys.stdin:
0078         if len(line) == 0:
0079             continue
0080         if end_reached:
0081             warn("Ignoring lines after END_EVENT_LISTING was reached")
0082             continue
0083         if not have_first_line:
0084             if not 'HepMC::Version' in line:
0085                 raise InvalidHepmc3Error('Not a valid HepMC3 file header:', line)
0086             have_first_line = True
0087             buffer.append(line)
0088         elif not have_second_line:
0089             if not 'HepMC::Asciiv3-START_EVENT_LISTING\n' == line:
0090                 raise InvalidHepmc3Error('Not a valid HepMC3 file header:', line)
0091             have_second_line = True
0092             buffer.append(line)
0093         elif 'HepMC::Asciiv3-END_EVENT_LISTING\n' == line:
0094             end_reached = True
0095             buffer.append(line)
0096         elif line[0] in ['A','W','T','N']:
0097             buffer.append(line)
0098         elif line[0] == 'E':
0099             flush_buffer(header, buffer)
0100             header = EventHeader(line)
0101             buffer = []
0102         else:
0103             if header is None:
0104                 raise InvalidHepmc3Error('Encountered invalid field before the first Event header:', line)
0105             if line[0] == 'V':
0106                 header.process_vertex(line)
0107             elif line[0] == 'P':
0108                 header.process_particle(line)
0109             elif line[0] not in ['W', 'A', 'U']:
0110                 warn('Ignoring unknown field:', line)
0111                 continue
0112             buffer.append(line)
0113     if not end_reached:
0114         warn("File does not end with END_EVENT_LISTING, appending")
0115         buffer.append('HepMC::Asciiv3-END_EVENT_LISTING\n')
0116     # final buffer flush at the end
0117     flush_buffer(header, buffer)