Update to fluka reader

Generators/FlukaReader/python/FlukaReaderAlg.py

@@ -12,7 +12,7 @@ import numpy as np
 import math
 
 class FlukaReader(EvgenAlg):
-    def __init__(self, name="FlukaReader", MCEventKey="BeamTruthEvent", file_name = "", dist = 0, randomSeed = None, nsamples = 1, test = False):
+    def __init__(self, name="FlukaReader", MCEventKey="BeamTruthEvent", file_name = "", dist = 0, randomSeed = None, nsamples = 1, test = True):
         super(FlukaReader,self).__init__(name=name)
         self.McEventKey = MCEventKey
         self.file_name = file_name
@@ -44,6 +44,7 @@ class FlukaReader(EvgenAlg):
         
         self.file.close()
         return StatusCode.Success
+    
 
     def fillEvent(self, evt):
         "This is called for every real event * the number of samplings"
@@ -81,6 +82,26 @@ class FlukaReader(EvgenAlg):
         "Convert cos(theta) wrt x or y axis to theta wrt z axis"
         return np.pi/2 - np.arccos(cosTheta)
 
+#     def theta_and_phi(self, thetaX, thetaY, twopi = True):
+#         "Converted thetaX and thetaY wrt Z to theta and phi"
+# 
+#         theta = thetaY
+# 
+# 
+#         phi = np.arctan(np.abs(thetaY)/np.abs(thetaX))
+# 
+#         if thetaX < 0 and thetaY > 0:
+#             phi = phi + np.pi/2.
+#         elif thetaX > 0 and thetaY < 0:
+#             phi = -phi
+#         elif  thetaX < 0 and thetaY < 0:
+#             phi = -phi - np.pi/2.
+# 
+#         if twopi and phi < 0:
+#             phi = 2*np.pi + phi
+
+        return theta, phi
+
     def pid(self, ftype):
         "Convert fluka particle type to PID"
         if ftype == 10:  # mu+
@@ -91,7 +112,7 @@ class FlukaReader(EvgenAlg):
             return 0
 
     def path_length(self, z, cosThetaX, cosThetaY):
-        "Get path length traversed in the material, taking into account incident angle"
+        "Get path length traversed in the material, taking into account incident angles"
         
         # Convert theta wrt x and y axis to wrt z axis
         thetaX = self.angle(cosThetaX)
@@ -268,11 +289,15 @@ class FlukaReader(EvgenAlg):
         e = entry["E"] * 1000.
         p = np.sqrt(e**2 - m**2)
 
-        thetaX = self.angle(entry["cosX"])
-        thetaY = self.angle(entry["cosY"])        
-        px = p * np.sin(thetaX) * np.cos(thetaY)
-        py = p * np.cos(thetaX) * np.sin(thetaY)        
-        pz = p * np.cos(thetaX) * np.cos(thetaY)
+        thetaX = self.angle(entry["cosX"])        
+        thetaY = self.angle(entry["cosY"])
+        #theta, phi = self.theta_and_phi(thetaX, thetaY)
+        theta = thetaY
+        phi = np.arctan(entry["cosY"]/entry["cosX"])
+                    
+        px = p * np.sin(theta) * np.cos(phi)
+        py = p * np.sin(theta) * np.sin(phi)        
+        pz = p * np.cos(theta) 
         
         mom = HepMC.FourVector(px, py, pz, e)
 
@@ -296,43 +321,67 @@ class FlukaReader(EvgenAlg):
         
         plt.figure()
         ebins = np.linspace(0, 5000, 50)
-        plt.hist(self.before["E"], bins=ebins, histtype='step', color = "g", fill = False)
-        plt.hist(self.after["E"], bins = ebins, histtype='step', color = "r", fill = False)
+        plt.xlabel("Energy")
+        plt.hist(self.before["E"], bins=ebins, histtype='step', color = "g", fill = False, label = "before")
+        plt.hist(self.after["E"], bins = ebins, histtype='step', color = "r", fill = False, label = "after")
         plt.gca().set_yscale('log')
+        plt.legend()
         plt.savefig("energy.eps")
 
         plt.figure()
+        plt.xlabel("Angle to beam in X dir")        
         thetaX =  np.pi/2. - np.arccos(np.array(self.before["cosX"]))
         thetaXout =  np.pi/2. - np.arccos(np.array(self.after["cosX"]))        
         tbins = np.linspace(-0.5, 0.5, 100)
-        plt.hist(thetaX, bins = tbins, histtype='step', color = "g", fill = False)
-        plt.hist(thetaXout, bins = tbins, histtype='step', color = "r", fill = False)
+        plt.hist(thetaX, bins = tbins, histtype='step', color = "g", fill = False, label = "before")
+        plt.hist(thetaXout, bins = tbins, histtype='step', color = "r", fill = False, label = "after")
         plt.gca().set_yscale('log')
+        plt.legend()
         plt.savefig("thetaX.eps")
 
         plt.figure()
+        plt.xlabel("Angle to beam in Y dir")        
         thetaY =  np.pi/2. - np.arccos(np.array(self.before["cosY"]))
         thetaYout =  np.pi/2. - np.arccos(np.array(self.after["cosY"]))        
-        plt.hist(thetaY, bins = tbins, histtype='step', color = "g", fill = False)
-        plt.hist(thetaYout, bins = tbins, histtype='step', color = "r", fill = False)
+        plt.hist(thetaY, bins = tbins, histtype='step', color = "g", fill = False, label = "before")
+        plt.hist(thetaYout, bins = tbins, histtype='step', color = "r", fill = False, label = "after")
         plt.gca().set_yscale('log')
+        plt.legend()
         plt.savefig("thetaY.eps")
 
         plt.figure()
+        plt.xlabel("Dispacement in X dir")        
         xbins = np.linspace(-300, 300, 100)
-        plt.hist(self.before["x"], bins = xbins, histtype='step', color = "g", fill = False)
-        plt.hist(self.after["x"], bins = xbins, histtype='step', color = "r", fill = False)
+        plt.hist(self.before["x"], bins = xbins, histtype='step', color = "g", fill = False, label = "before")
+        plt.hist(self.after["x"], bins = xbins, histtype='step', color = "r", fill = False, label = "after")
         plt.gca().set_yscale('log')
+        plt.legend()
         plt.savefig("x.eps")
 
-        plt.figure()
-        plt.hist(self.before["y"], bins = xbins, histtype='step', color = "g", fill = False)
-        plt.hist(self.after["y"], bins = xbins, histtype='step', color = "r", fill = False)        
+        plt.figure() 
+        plt.xlabel("Dispacement in Y dir")               
+        plt.hist(self.before["y"], bins = xbins, histtype='step', color = "g", fill = False, label = "before")
+        plt.hist(self.after["y"], bins = xbins, histtype='step', color = "r", fill = False, label = "after")        
         plt.gca().set_yscale('log')
+        plt.legend()
         plt.savefig("y.eps")
 
         return
 
+def getNEvents(fname, maxEvents):
+    "Work out how many events are in the file"
+
+    n = 0
+    with open(fname) as f:
+        n = sum(1 for _ in f)
+
+    if maxEvents != -1 and n > maxEvents:
+        n = maxEvents
+
+    print(">>> Setting number of real events to", n)
+
+    return n
+
 if __name__ == "__main__":
 
 #     from AthenaCommon.AlgSequence import AlgSequence
@@ -391,5 +440,5 @@ if __name__ == "__main__":
     from OutputStreamAthenaPool.OutputStreamConfig import OutputStreamCfg
     cfg.merge(OutputStreamCfg(ConfigFlags, "EVNT", itemList, disableEventTag = True))
 
-    sc = cfg.run(maxEvents = args.nevents * args.nsamples)
+    sc = cfg.run(maxEvents = getNEvents(args.file, args.nevents) * args.nsamples)
     sys.exit(not sc.isSuccess())