diff --git a/src/sntools/event.py b/src/sntools/event.py index 8b0ef02..6b759c4 100644 --- a/src/sntools/event.py +++ b/src/sntools/event.py @@ -1,3 +1,5 @@ +import uproot + class Event(object): """A single neutrino interaction in the detector.""" @@ -16,26 +18,62 @@ def __setattr__(self, name, value): raise AttributeError(f"{name} is a list. Append to it instead of overwriting it.") object.__setattr__(self, name, value) - def nuance_string(self, i): + +class EventWriter: + def __init__(self, format, outfile): + self.format = format + + file_extensions = {"NUANCE": ".kin", "RATPAC": ".txt", "ROOT_JUNO": ".root"} + if not outfile.endswith(file_extensions[format]): + outfile += file_extensions[format] + + if format in ('NUANCE', 'RATPAC'): + self.outfile = open(outfile, "w") + elif format == 'ROOT_JUNO': + self.outfile = uproot.recreate(outfile) + + def write_preamble(self, contents: str): + """Write preamble to output file. Not supported by all formats.""" + if self.format in ('NUANCE', 'RATPAC'): + for line in contents.splitlines(): + self.outfile.write(f"# {line}\n") + + def write_events(self, events): + """Write list of Event objects to output file.""" + if self.format == 'NUANCE': + self._write_nuance_events(events) + elif self.format == 'RATPAC': + self._write_ratpac_events(events) + elif self.format == 'ROOT_JUNO': + self._write_juno_events(events) + + self.outfile.close() + + def _write_nuance_events(self, events): """Return NUANCE-formatted representation of event for writing to output file. Input: + evt: Event object i: number of event Output: String describing event.""" - s = "$ begin\n" - s += f"$ nuance {self.code}\n" - s += f"$ vertex {self.vertex[0]:.5f} {self.vertex[1]:.5f} {self.vertex[2]:.5f} {self.time:.8f}\n" - for (pid, e, dirx, diry, dirz) in self.incoming_particles: - s += f"$ track {pid} {e:.5f} {dirx:.5f} {diry:.5f} {dirz:.5f} -1\n" - s += f"$ info 0 0 {i}\n" - for (pid, e, dirx, diry, dirz) in self.outgoing_particles: - s += f"$ track {pid} {e:.5f} {dirx:.5f} {diry:.5f} {dirz:.5f} 0\n" - s += "$ end\n" - return s - - def ratpac_string(self, i, events): + for (i, evt) in enumerate(events): + s = "$ begin\n" + s += f"$ nuance {evt.code}\n" + s += f"$ vertex {evt.vertex[0]:.5f} {evt.vertex[1]:.5f} {evt.vertex[2]:.5f} {evt.time:.8f}\n" + for (pid, e, dirx, diry, dirz) in evt.incoming_particles: + s += f"$ track {pid} {e:.5f} {dirx:.5f} {diry:.5f} {dirz:.5f} -1\n" + s += f"$ info 0 0 {i}\n" + for (pid, e, dirx, diry, dirz) in evt.outgoing_particles: + s += f"$ track {pid} {e:.5f} {dirx:.5f} {diry:.5f} {dirz:.5f} 0\n" + s += "$ end\n" + + self.outfile.write(s) + + self.outfile.write("$ stop\n") + + def _write_ratpac_events(self, events): """Return RAT-PAC readable HEPEVT-style representation of event for writing to output file. Input: @@ -48,35 +86,34 @@ def ratpac_string(self, i, events): mm = 10 # convert from cm ns = 1000000 # convert from ms - dt = self.time - if i > 0: - dt -= events[i - 1].time - - s = f"{len(self.outgoing_particles)}\n" - for idx, (pid, e, dirx, diry, dirz) in enumerate(self.outgoing_particles): - mass = 0.0 - if pid == 11 or pid == -11: - mass = 0.5109907 - if pid == 22: - mass = 0.0 - if pid == 2112: - mass = 939.56563 - if pid == 2212: - mass = 938.27205 - p2 = (e**2) - (mass**2) - p = p2**0.5 - px = dirx * p - py = diry * p - pz = dirz * p - if idx > 0: - dt = 0.0 - s += f"1 {pid} 0 0 {px * GeV:.8e} {py * GeV:.8e} {pz * GeV:.8e} {mass * GeV:.8e} {dt * ns:.5e} {self.vertex[0] * mm:.5e} {self.vertex[1] * mm:.5e} {self.vertex[2] * mm:.5e}\n" - return s - - def juno_string(self, i, outfile): + for (i, evt) in enumerate(events): + dt = evt.time + if i > 0: + dt -= events[i - 1].time + + s = f"{len(evt.outgoing_particles)}\n" + for idx, (pid, e, dirx, diry, dirz) in enumerate(evt.outgoing_particles): + if pid == 11 or pid == -11: + mass = 0.5109907 + elif pid == 2112: + mass = 939.56563 + elif pid == 2212: + mass = 938.27205 + else: + mass = 0.0 + p = (e**2 - mass**2)**0.5 + px = dirx * p + py = diry * p + pz = dirz * p + if idx > 0: + dt = 0.0 + s += f"1 {pid} 0 0 {px * GeV:.8e} {py * GeV:.8e} {pz * GeV:.8e} {mass * GeV:.8e} {dt * ns:.5e} {evt.vertex[0] * mm:.5e} {evt.vertex[1] * mm:.5e} {evt.vertex[2] * mm:.5e}\n" + + self.outfile.write(s) + + def _write_juno_events(self, events): class EVENT(): - def __init__(self): self.nparticles = 0 self.t = [0,0] @@ -88,44 +125,42 @@ def __init__(self): self.pdgid = [0,0] self.origPDGID = 0 self.channel = 0 - - def fill_root(self,outfile): - outfile["SNEvents"].extend({"pdgid": [self.pdgid],"px":[self.px],"py":[self.py],"pz":[self.pz],"t":[self.t],"m":[self.m], - "nuE":[self.nuE], "nparticles":[self.nparticles], "origPDGID":[self.origPDGID], "channel":[self.channel]}) - - evt = EVENT() - for idx, (pid, e, dirx, diry, dirz) in enumerate(self.outgoing_particles): - mass = 0.0 - if pid == 11 or pid == -11: - mass = 0.5109907 - if pid == 22: - mass = 0.0 - if pid == 2112: - mass = 939.56563 - if pid == 2212: - mass = 938.27205 - p2 = (e**2) - (mass**2) - p = p2**0.5 - evt.px[idx] = dirx * p - evt.py[idx] = diry * p - evt.pz[idx] = dirz * p - evt.m[idx] = mass - evt.pdgid[idx]=pid - - if len(self.outgoing_particles) <2: - #is elastic scattering, second particle is a neutrino, not visible - evt.px[1]=0 - evt.py[1]=0 - evt.pz[1]=0 - evt.m[1]=0 - evt.pdgid[1]=0 - - evt.nparticles = len(self.outgoing_particles) - evt.nuE = self.incoming_particles[0][1] - evt.t = [self.time*1e6,0] - evt.origPDGID = self.incoming_particles[0][0] - evt.channel = self.code - - EVENT.fill_root(evt,outfile) - + self.outfile.mktree("SNEvents",{"nparticles": "uint64", "origPDGID":"int32", "nuE":"double", "pdgid": ("int32",(2,)),"t": ("float64",(2,)), + "px": ("float64",(2,)),"py":("float64",(2,)),"pz":("float64",(2,)),"m":("float64",(2,)), "channel": "int64"}) + + for orig_evt in events: + evt = EVENT() + + for idx, (pid, e, dirx, diry, dirz) in enumerate(orig_evt.outgoing_particles): + if pid == 11 or pid == -11: + mass = 0.5109907 + elif pid == 2112: + mass = 939.56563 + elif pid == 2212: + mass = 938.27205 + else: + mass = 0.0 + p = (e**2 - mass**2)**0.5 + evt.px[idx] = dirx * p + evt.py[idx] = diry * p + evt.pz[idx] = dirz * p + evt.m[idx] = mass + evt.pdgid[idx] = pid + + if len(orig_evt.outgoing_particles) < 2: + # is elastic scattering, second particle is a neutrino, not visible + evt.px[1] = 0 + evt.py[1] = 0 + evt.pz[1] = 0 + evt.m[1] = 0 + evt.pdgid[1] = 0 + + evt.nparticles = len(orig_evt.outgoing_particles) + evt.nuE = orig_evt.incoming_particles[0][1] + evt.t = [orig_evt.time*1e6, 0] + evt.origPDGID = orig_evt.incoming_particles[0][0] + evt.channel = orig_evt.code + + self.outfile["SNEvents"].extend({"pdgid": [evt.pdgid],"px":[evt.px],"py":[evt.py],"pz":[evt.pz],"t":[evt.t],"m":[evt.m], + "nuE":[evt.nuE], "nparticles":[evt.nparticles], "origPDGID":[evt.origPDGID], "channel":[evt.channel]}) diff --git a/src/sntools/genevts.py b/src/sntools/genevts.py index b682d24..757bc43 100644 --- a/src/sntools/genevts.py +++ b/src/sntools/genevts.py @@ -5,7 +5,6 @@ from datetime import datetime from importlib import import_module import random -import uproot try: import sntools # if sntools was installed via pip @@ -20,6 +19,7 @@ from sntools.channel import gen_evts from sntools.detectors import Detector, supported_detectors +from sntools.event import EventWriter from sntools.formats import CompositeFlux, SNEWPYCompositeFlux from sntools.transformation import Transformation, SNEWPYTransformation @@ -74,25 +74,11 @@ def main(): for evt in events: evt.vertex = args.detector.generate_random_vertex() - if args.mcformat in ('NUANCE', 'RATPAC'): - with open(args.output, "w") as outfile: - if args.verbose: # write parameters to file as a comment - outfile.write(f"# Generated on {datetime.now()} with the options:\n") - outfile.write(f"# {args}\n") - if args.mcformat == 'NUANCE': - for (i, evt) in enumerate(events): - outfile.write(evt.nuance_string(i)) - outfile.write("$ stop\n") - if args.mcformat == 'RATPAC': - for (i, evt) in enumerate(events): - outfile.write(evt.ratpac_string(i, events)) - if args.mcformat == 'ROOT_JUNO': - fname = args.output+".root" - root_outfile = uproot.recreate(fname) - root_outfile.mktree("SNEvents",{"nparticles": "uint64", "origPDGID":"int32", "nuE":"double", "pdgid": ("int32",(2,)),"t": ("float64",(2,)), - "px": ("float64",(2,)),"py":("float64",(2,)),"pz":("float64",(2,)),"m":("float64",(2,)), "channel": "int64"}) - for (i, evt) in enumerate(events): - evt.juno_string(i, root_outfile) + writer = EventWriter(args.mcformat, args.output) + if args.verbose: # write parameters to file as a comment + writer.write_preamble(f"Generated on {datetime.now()} with the options:\n{args}") + writer.write_events(events) + def parse_command_line_options(): """Define and parse command line options.""" @@ -114,7 +100,7 @@ def parse_command_line_options(): parser.add_argument("-f", "--format", metavar="FORMAT", choices=choices, default=choices[1], help="Format of input file(s). Choices: %(choices)s. Default: %(default)s.") - parser.add_argument("-o", "--output", metavar="FILE", default="outfile.kin", help="Name of the output file. Default: %(default)s.") + parser.add_argument("-o", "--output", metavar="FILE", default="outfile", help="Name of the output file. Default: %(default)s. (File extension is added automatically based on output format.)") choices = ("NUANCE", "RATPAC","ROOT_JUNO") parser.add_argument("-m", "--mcformat", metavar="MCFORMAT", choices=choices, default=choices[0],