From 8cb45837c89cbad24cae41a8c1c7690f924505c1 Mon Sep 17 00:00:00 2001 From: Jost Migenda Date: Thu, 19 Mar 2026 12:35:02 +0000 Subject: [PATCH 1/5] Move output file writing to dedicated class --- src/sntools/event.py | 76 +++++++++++++++++++++++++++++++----------- src/sntools/genevts.py | 26 ++++----------- 2 files changed, 63 insertions(+), 39 deletions(-) diff --git a/src/sntools/event.py b/src/sntools/event.py index 8b0ef02..0dd7b14 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,63 @@ 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 + + # TODO: check file endings for consistency with format, e.g. .root for ROOT_JUNO + + 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.""" + match self.format: + case 'NUANCE': + for (i, evt) in enumerate(events): + self.outfile.write(self._nuance_string(evt, i)) + self.outfile.write("$ stop\n") + case 'RATPAC': + for (i, evt) in enumerate(events): + self.outfile.write(self._ratpac_string(evt, i, events)) + case 'ROOT_JUNO': + 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 (i, evt) in enumerate(events): + self._juno_string(evt, i, self.outfile) + + self.outfile.close() + + def _nuance_string(self, evt, i): """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"$ 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 self.outgoing_particles: + 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" return s - def ratpac_string(self, i, events): + def _ratpac_string(self, evt, i, events): """Return RAT-PAC readable HEPEVT-style representation of event for writing to output file. Input: @@ -48,12 +87,12 @@ def ratpac_string(self, i, events): mm = 10 # convert from cm ns = 1000000 # convert from ms - dt = self.time + dt = evt.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): + s = f"{len(evt.outgoing_particles)}\n" + for idx, (pid, e, dirx, diry, dirz) in enumerate(evt.outgoing_particles): mass = 0.0 if pid == 11 or pid == -11: mass = 0.5109907 @@ -70,10 +109,10 @@ def ratpac_string(self, i, events): 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" + 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" return s - def juno_string(self, i, outfile): + def _juno_string(self, orig_evt, i, outfile): class EVENT(): @@ -95,7 +134,7 @@ def fill_root(self,outfile): "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): + for idx, (pid, e, dirx, diry, dirz) in enumerate(orig_evt.outgoing_particles): mass = 0.0 if pid == 11 or pid == -11: mass = 0.5109907 @@ -113,7 +152,7 @@ def fill_root(self,outfile): evt.m[idx] = mass evt.pdgid[idx]=pid - if len(self.outgoing_particles) <2: + 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 @@ -121,11 +160,10 @@ def fill_root(self,outfile): 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 + 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 EVENT.fill_root(evt,outfile) - diff --git a/src/sntools/genevts.py b/src/sntools/genevts.py index b682d24..ae7c733 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.""" From aa399278a60df2d06ac9bccf9dacfa46869e0a0a Mon Sep 17 00:00:00 2001 From: Jost Migenda Date: Thu, 19 Mar 2026 13:39:09 +0000 Subject: [PATCH 2/5] minor code cleanup --- src/sntools/event.py | 52 +++++++++++++++++++------------------------- 1 file changed, 22 insertions(+), 30 deletions(-) diff --git a/src/sntools/event.py b/src/sntools/event.py index 0dd7b14..60fc5cb 100644 --- a/src/sntools/event.py +++ b/src/sntools/event.py @@ -93,17 +93,15 @@ def _ratpac_string(self, evt, i, events): s = f"{len(evt.outgoing_particles)}\n" for idx, (pid, e, dirx, diry, dirz) in enumerate(evt.outgoing_particles): - mass = 0.0 if pid == 11 or pid == -11: mass = 0.5109907 - if pid == 22: - mass = 0.0 - if pid == 2112: + elif pid == 2112: mass = 939.56563 - if pid == 2212: + elif pid == 2212: mass = 938.27205 - p2 = (e**2) - (mass**2) - p = p2**0.5 + else: + mass = 0.0 + p = (e**2 - mass**2)**0.5 px = dirx * p py = diry * p pz = dirz * p @@ -115,7 +113,6 @@ def _ratpac_string(self, evt, i, events): def _juno_string(self, orig_evt, i, outfile): class EVENT(): - def __init__(self): self.nparticles = 0 self.t = [0,0] @@ -127,43 +124,38 @@ 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(orig_evt.outgoing_particles): - mass = 0.0 if pid == 11 or pid == -11: mass = 0.5109907 - if pid == 22: - mass = 0.0 - if pid == 2112: + elif pid == 2112: mass = 939.56563 - if pid == 2212: + elif pid == 2212: mass = 938.27205 - p2 = (e**2) - (mass**2) - p = p2**0.5 + 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 + 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 + 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.t = [orig_evt.time*1e6, 0] evt.origPDGID = orig_evt.incoming_particles[0][0] evt.channel = orig_evt.code - EVENT.fill_root(evt,outfile) + 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]}) From 3aece0b2ef83ab4adf758df466659e8731019354 Mon Sep 17 00:00:00 2001 From: Jost Migenda Date: Thu, 19 Mar 2026 13:45:38 +0000 Subject: [PATCH 3/5] unify _write_*_events functions for all formats --- src/sntools/event.py | 164 ++++++++++++++++++++++--------------------- 1 file changed, 84 insertions(+), 80 deletions(-) diff --git a/src/sntools/event.py b/src/sntools/event.py index 60fc5cb..753f475 100644 --- a/src/sntools/event.py +++ b/src/sntools/event.py @@ -40,21 +40,15 @@ def write_events(self, events): """Write list of Event objects to output file.""" match self.format: case 'NUANCE': - for (i, evt) in enumerate(events): - self.outfile.write(self._nuance_string(evt, i)) - self.outfile.write("$ stop\n") + self._write_nuance_events(events) case 'RATPAC': - for (i, evt) in enumerate(events): - self.outfile.write(self._ratpac_string(evt, i, events)) + self._write_ratpac_events(events) case 'ROOT_JUNO': - 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 (i, evt) in enumerate(events): - self._juno_string(evt, i, self.outfile) + self._write_juno_events(events) self.outfile.close() - def _nuance_string(self, evt, i): + def _write_nuance_events(self, events): """Return NUANCE-formatted representation of event for writing to output file. Input: @@ -63,18 +57,22 @@ def _nuance_string(self, evt, i): Output: String describing event.""" - 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" - return s - - def _ratpac_string(self, evt, 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: @@ -87,30 +85,32 @@ def _ratpac_string(self, evt, i, events): mm = 10 # convert from cm ns = 1000000 # convert from ms - 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" - return s - - def _juno_string(self, orig_evt, 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): @@ -125,37 +125,41 @@ def __init__(self): self.origPDGID = 0 self.channel = 0 - 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]}) + 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]}) From 0bafcd4bffd255231ff5c38b1537739de4578f21 Mon Sep 17 00:00:00 2001 From: Jost Migenda Date: Thu, 19 Mar 2026 16:57:15 +0000 Subject: [PATCH 4/5] Handle file extension in EventWriter --- src/sntools/event.py | 4 +++- src/sntools/genevts.py | 2 +- 2 files changed, 4 insertions(+), 2 deletions(-) diff --git a/src/sntools/event.py b/src/sntools/event.py index 753f475..87b3515 100644 --- a/src/sntools/event.py +++ b/src/sntools/event.py @@ -23,7 +23,9 @@ class EventWriter: def __init__(self, format, outfile): self.format = format - # TODO: check file endings for consistency with format, e.g. .root for ROOT_JUNO + 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") diff --git a/src/sntools/genevts.py b/src/sntools/genevts.py index ae7c733..757bc43 100644 --- a/src/sntools/genevts.py +++ b/src/sntools/genevts.py @@ -100,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], From 43e283f78fb358c30e06ebbff3599323a57f7635 Mon Sep 17 00:00:00 2001 From: Jost Migenda Date: Thu, 19 Mar 2026 16:59:00 +0000 Subject: [PATCH 5/5] restore Python 3.9 compatibility --- src/sntools/event.py | 13 ++++++------- 1 file changed, 6 insertions(+), 7 deletions(-) diff --git a/src/sntools/event.py b/src/sntools/event.py index 87b3515..6b759c4 100644 --- a/src/sntools/event.py +++ b/src/sntools/event.py @@ -40,13 +40,12 @@ def write_preamble(self, contents: str): def write_events(self, events): """Write list of Event objects to output file.""" - match self.format: - case 'NUANCE': - self._write_nuance_events(events) - case 'RATPAC': - self._write_ratpac_events(events) - case 'ROOT_JUNO': - self._write_juno_events(events) + 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()