Move output file writing to dedicated class

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]})

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],