ENH: Fit WEPL to angle standard deviation in pctweplfit

Author: acoussat

applications/pctweplfit/pctweplfit.py

@@ -116,6 +116,7 @@ def add_detector(name, translation, attach_to_phantom=False):
             'EventID',
             'TrackID',
             'Position',
+            'Direction',
             'PreGlobalTime',
             'KineticEnergy'
         ]
@@ -138,10 +139,11 @@ def add_detector(name, translation, attach_to_phantom=False):
     sim.run()
 
 
-def process_phantom_length(phantom_length, output, path_type, number_of_detectors, tofs, wepls, elosses, verbose):
+def process_phantom_length(phantom_length, output, path_type, number_of_detectors, tofs, wepls, elosses, angles, verbose):
     tofs_phantom_length = []
     wepls_phantom_length = []
     elosses_phantom_length = []
+    angles_phantom_length = []
 
     data = uproot.concatenate(f'{output}/l{int(phantom_length)}_*.root', library='np')
     pv(verbose, "Loaded", len(data['EventID']), "events for phantom length", phantom_length)
@@ -163,15 +165,20 @@ def process_phantom_length(phantom_length, output, path_type, number_of_detector
         if path_type != 'phantom_length' and number_of_hits < number_of_detectors + 2 and phantom_length > 0.:
             continue
 
+        us = data['Position_X'][event_mask]
+        vs = data['Position_Y'][event_mask]
+        ws = data['Position_Z'][event_mask]
+        dus = data['Direction_X'][event_mask]
+        dvs = data['Direction_Y'][event_mask]
+        dws = data['Direction_Z'][event_mask]
+
         times = data['PreGlobalTime'][event_mask]
+
         tof = times[-1] - times[0]
 
         if phantom_length == 0.:
             wepl = 0.
         else:
-            us = data['Position_X'][event_mask]
-            vs = data['Position_Y'][event_mask]
-            ws = data['Position_Z'][event_mask]
             if path_type == 'phantom_length':
                 # Length of the phantom
                 wepl = phantom_length
@@ -189,13 +196,24 @@ def process_phantom_length(phantom_length, output, path_type, number_of_detector
 
         eloss = data['KineticEnergy'][event_mask][0] - data['KineticEnergy'][event_mask][-1]
 
+        d_in_uw = np.array([dus[0],dws[0]])
+        d_in_vw = np.array([dvs[0],dws[0]])
+        d_out_uw = np.array([dus[-1],dws[-1]])
+        d_out_vw = np.array([dvs[-1],dws[-1]])
+        angle = [
+            np.arccos(np.minimum(1., d_in_uw @ d_out_uw / (np.sqrt(d_in_uw @ d_in_uw) * np.sqrt(d_out_uw @ d_out_uw)))),
+            np.arccos(np.minimum(1., d_in_vw @ d_out_vw / (np.sqrt(d_in_vw @ d_in_vw) * np.sqrt(d_out_vw @ d_out_vw))))
+        ]
+
         tofs_phantom_length.append(tof)
         wepls_phantom_length.append(wepl)
         elosses_phantom_length.append(eloss)
+        angles_phantom_length += angle
 
     tofs[phantom_length] = tofs_phantom_length
     wepls[phantom_length] = wepls_phantom_length
     elosses[phantom_length] = elosses_phantom_length
+    angles[phantom_length] = angles_phantom_length
 
 
 def tof_fit(
@@ -214,20 +232,21 @@ def tof_fit(
     tofs = manager.dict()
     wepls = manager.dict()
     elosses = manager.dict()
+    angles = manager.dict()
 
     results = []
     with Pool() as pool:
         for phantom_length in phantom_lengths:
-            result = pool.apply_async(process_phantom_length, (phantom_length, output, path_type, number_of_detectors, tofs, wepls, elosses, verbose))
+            result = pool.apply_async(process_phantom_length, (phantom_length, output, path_type, number_of_detectors, tofs, wepls, elosses, angles, verbose))
             results.append(result)
         pool.close()
         pool.join()
     for result in results:
         result.get()
 
-    def fit(xs, ys, xlabel, ylabel):
+    def fit(xs, ys, xlabel, ylabel, xreduction=np.median):
 
-        xmedians = [np.median(xs[phantom_length]) for phantom_length in phantom_lengths]
+        xmedians = [xreduction(xs[phantom_length]) for phantom_length in phantom_lengths]
         ymedians = [np.median(ys[phantom_length]) for phantom_length in phantom_lengths]
         xpercentile25 = [np.percentile(xs[phantom_length], 25) for phantom_length in phantom_lengths]
         xpercentile75 = [np.percentile(xs[phantom_length], 75) for phantom_length in phantom_lengths]
@@ -248,15 +267,14 @@ def fit(xs, ys, xlabel, ylabel):
 
         if display or savefig:
             plt.figure()
-            plt.plot(xmedians, ymedians, '+', label="Medians")
+            plt.plot(xmedians, ymedians, '+')
 
             tof_xs = np.linspace(np.min(xmedians), np.max(xmedians), 100)
             for d, p in ps.items():
-                plt.plot(tof_xs, np.polyval(p, tof_xs), label=f"Polynomial fit (degree {d})")
+                plt.plot(tof_xs, np.polyval(p, tof_xs))
 
             plt.xlabel(xlabel)
             plt.ylabel(ylabel)
-            plt.legend()
 
             if savefig:
                 plt.savefig(f'{output}/{xlabel}_to_{ylabel}_fit.pdf')
@@ -269,6 +287,7 @@ def fit(xs, ys, xlabel, ylabel):
 
     fit(tofs, wepls, 'tof', 'wepl')
     fit(elosses, wepls, 'eloss', 'wepl')
+    fit(angles, wepls, xlabel='angle', ylabel='wepl', xreduction=np.std)
 
 
 def pctweplfit(