ENH: add track angle histogram #46

testbeam_analysis/result_analysis.py

@@ -2,6 +2,7 @@
 from __future__ import division
 
 import logging
+import progressbar
 import re
 from collections import Iterable
 
@@ -847,3 +848,56 @@ def calculate_efficiency(input_tracks_file, input_alignment_file, output_pdf, bi
                         out_pass[:] = total_track_density_with_DUT_hit.T
                         out_total[:] = total_track_density.T
     return efficiencies, pass_tracks, total_tracks
+
+
+def track_angle(input_tracks_file, output_track_angle_file, use_duts=None, output_pdf=None):
+    '''Calculates and histograms the track angle of the fitted tracks for selected DUTs.
+    Parameters
+    ----------
+    input_tracks_file : string
+        file name with the tracks table
+    output_track_angle_file: string
+        output pytables file with fitted means and sigmas of track angles for selected DUTs
+    use_duts : iterable
+        The duts to plot/calculate the track angle. If None all duts in the input_tracks_file are used
+    output_pdf : string
+        file name for the output pdf
+    '''
+    logging.info('=== Calculate track angles ===')
+
+    slopes = {}
+    with tb.open_file(input_tracks_file, 'r') as in_file_h5:
+        n_duts = len(in_file_h5.list_nodes("/"))  # determine number of DUTs
+
+        if use_duts is None:
+            use_duts = range(n_duts)
+        else:
+            use_duts = use_duts
+        dut_index = [('Tracks_DUT_%i' % i) for i in use_duts]
+
+        widgets = ['', progressbar.Percentage(), ' ',
+                   progressbar.Bar(marker='*', left='|', right='|'),
+                   ' ', progressbar.AdaptiveETA()]
+        progress_bar = progressbar.ProgressBar(widgets=widgets,
+                                               maxval=len(use_duts),
+                                               term_width=80)
+        progress_bar.start()
+
+        i = 0  # counting number of progressed DUTs, needed for progressbar
+        for node in in_file_h5.root:  # loop through all DUTs in track table
+            if node.name in dut_index:  # only store track slopes of selected DUTs
+                slopes.update({node.name: [node[::10]['slope_0'],
+                                           node[::10]['slope_1']]})
+                i = i + 1
+                progress_bar.update(i)
+        progress_bar.finish()
+
+    result = np.zeros(shape=(n_duts,), dtype=[('mean_slope_0', np.float), ('mean_slope_1', np.float), ('sigma_slope_0', np.float), ('sigma_slope_1', np.float)])
+
+    logging.info('=== Plot track angles ===')
+
+    plot_utils.plot_track_angle(input_track_slopes=slopes, output_data=result, output_pdf=output_pdf)
+
+    with tb.open_file(output_track_angle_file, mode="w") as out_file_h5:
+        result_table = out_file_h5.create_table(out_file_h5.root, name='TrackAngle', description=result.dtype, title='Fitted means and sigmas for track angles', filters=tb.Filters(complib='blosc', complevel=5, fletcher32=False))
+        result_table.append(result)

testbeam_analysis/tools/plot_utils.py

@@ -1067,3 +1067,37 @@ def efficiency_plots(hit_hist, track_density, track_density_with_DUT_hit, effici
         output_fig.savefig()
     else:
         logging.warning('Cannot create efficiency plots, since all pixels are masked')
+
+
+def plot_track_angle(input_track_slopes, output_data, output_pdf):
+    direction = ['x', 'y', 'z']
+
+    with PdfPages(output_pdf) as output_fig:
+        i = 0  # counting number of progressed DUTs
+        for key in input_track_slopes:
+            for j in (0, 1):  # iterate over slopes
+                min = np.min(input_track_slopes[key][j])
+                max = np.max(input_track_slopes[key][j])
+                edges = np.arange(min - 0.00005, max + 0.00005, 0.0001)
+                bin_center = (edges[1:] + edges[:-1]) / 2.0
+
+                sigma = np.std(np.arctan(input_track_slopes[key][j]))  # initial guess for sigma
+                mean = np.mean(np.arctan(input_track_slopes[key][j]))  # initial guess for mean
+
+                histo = plt.hist(np.arctan(input_track_slopes[key][j]), edges, label='Angular Distribution', color='#4361ba')
+                fit, cov = curve_fit(testbeam_analysis.tools.analysis_utils.gauss, bin_center, histo[0], p0=[np.amax(histo[0]), mean, sigma])
+
+                x_gauss = np.arange(np.min(edges), np.max(edges), step=0.00001)
+                plt.plot(x_gauss, testbeam_analysis.tools.analysis_utils.gauss(x_gauss, *fit), color='r', label='Gauss-Fit:\nMean: %.5f,\nSigma: %.5f' % (fit[1], fit[2]))
+
+                output_data['mean_slope_%i' % j][i] = fit[1]
+                output_data['sigma_slope_%i' % j][i] = fit[2]
+
+                plt.title('Angular Distribution of Fitted Tracks for DUT%i' % i)
+                plt.ylabel('#')
+                plt.xlabel('Track angle in %s-direction / rad' % direction[j])
+                plt.legend(loc='best', fancybox=True, frameon=True)
+
+                output_fig.savefig()
+                plt.close()
+            i = i + 1