From 97f7b0ba8c4e23ff3e22c648efa55261cbeaf48d Mon Sep 17 00:00:00 2001
From: Eric Teunis de Boone <ericteunis@deboone.nl>
Date: Mon, 27 Mar 2023 16:54:47 +0200
Subject: [PATCH] ZH: show fitting information in phase plots

---
 .../bb_measure_clock_phase.py                 |  4 ++
 .../airshower_beacon_simulation/lib/figlib.py | 56 ++++++++++---------
 2 files changed, 35 insertions(+), 25 deletions(-)

diff --git a/simulations/airshower_beacon_simulation/bb_measure_clock_phase.py b/simulations/airshower_beacon_simulation/bb_measure_clock_phase.py
index 1b66d77..2dfb1fe 100755
--- a/simulations/airshower_beacon_simulation/bb_measure_clock_phase.py
+++ b/simulations/airshower_beacon_simulation/bb_measure_clock_phase.py
@@ -33,6 +33,7 @@ if __name__ == "__main__":
     ####
     fname_dir = args.data_dir
     antennas_fname = path.join(fname_dir, beacon.antennas_fname)
+    snr_fname = path.join(fname_dir, beacon.snr_fname)
 
     fig_dir = args.fig_dir # set None to disable saving
 
@@ -156,6 +157,8 @@ if __name__ == "__main__":
                 ##
                 ## Histogram
                 ##
+                snrs = beacon.read_snr_file(snr_fname)
+
                 fig = figlib.phase_comparison_figure(
                         loc_c,
                         actual_clock_phases,
@@ -163,6 +166,7 @@ if __name__ == "__main__":
                         f_beacon=f_beacon,
                         figsize=figsize,
                         fit_gaussian=plot_residuals,
+                        mean_snr=snrs['mean']
                         )
 
                 if plot_residuals:
diff --git a/simulations/airshower_beacon_simulation/lib/figlib.py b/simulations/airshower_beacon_simulation/lib/figlib.py
index 86edd8b..36f1c61 100644
--- a/simulations/airshower_beacon_simulation/lib/figlib.py
+++ b/simulations/airshower_beacon_simulation/lib/figlib.py
@@ -11,10 +11,13 @@ def phase_comparison_figure(
         f_beacon=None,
         hist_kwargs={},
         sc_kwargs={},
+        text_kwargs={},
         colors=['blue', 'orange'],
         legend_on_scatter=True,
         secondary_axis='time',
         fit_gaussian=False,
+        mean_snr=None,
+        return_fit_info=False,
         **fig_kwargs
         ):
     """
@@ -22,7 +25,14 @@ def phase_comparison_figure(
     by both plotting the values, and the residuals.
     """
     default_fig_kwargs = dict(sharex=True)
+    default_hist_kwargs = dict(bins='sqrt', density=False, alpha=0.8, histtype='step')
+    default_text_kwargs = dict(fontsize=14, verticalalignment='top')
+    default_sc_kwargs = dict(alpha=0.6, ls='none')
+
     fig_kwargs = {**default_fig_kwargs, **fig_kwargs}
+    hist_kwargs = {**default_hist_kwargs, **hist_kwargs}
+    text_kwargs = {**default_text_kwargs, **text_kwargs}
+    sc_kwargs = {**default_sc_kwargs, **sc_kwargs}
 
     do_hist_plot = hist_kwargs is not False
     do_scatter_plot = sc_kwargs is not False
@@ -46,41 +56,34 @@ def phase_comparison_figure(
         secax = axs[0].secondary_xaxis('top', functions=functions)
 
     # Histogram
+    fit_info = {}
     if do_hist_plot:
         i=0
-        default_hist_kwargs = dict(bins='sqrt', density=False, alpha=0.8, histtype='step')
-        hist_kwargs = {**default_hist_kwargs, **hist_kwargs}
 
         axs[i].set_ylabel("#")
 
-        _counts, _bins, _patches = axs[i].hist(measured_phases, color=colors[0], label='Measured', ls='solid', **hist_kwargs)
+        this_kwargs = dict(
+                ax = axs[i],
+                text_kwargs=text_kwargs,
+                hist_kwargs={**hist_kwargs, **dict(label='Measured', color=colors[0], ls='solid')},
+                mean_snr=mean_snr,
+                )
+
+        if fit_gaussian:
+            this_kwargs['fit_distr'] = 'gaussian'
+
+        _, fit_info = fitted_histogram_figure(
+                measured_phases,
+                **this_kwargs
+                )
+
         if not plot_residuals: # also plot the true clock phases
+            _bins = fit_info['bins']
             axs[i].hist(true_phases, color=colors[1], label='Actual', ls='dashed', **{**hist_kwargs, **dict(bins=_bins)})
 
-        if fit_gaussian:# Fit a gaussian to the histogram
-            gauss_kwargs = dict(color=colors[0], ls='dotted')
-            xs = np.linspace(min(_bins), max(_bins))
-
-            mean, std = norm.fit(measured_phases)
-            dx = _bins[1] - _bins[0]
-            scale = len(measured_phases)*dx
-            fit_ys = norm.pdf(xs, mean, std) * scale
-
-            axs[i].axvline(mean, ls='dotted', color='k', lw=2)
-
-            axs[i].plot( xs, fit_ys, label='fit', **gauss_kwargs)
-
-            # put mean, sigma and chisq on plot
-            text_str = f"$\\mu$ = {mean: .2e}\n$\\sigma$ = {std: .2e}"
-            text_kwargs = dict( transform=axs[i].transAxes, fontsize=14, verticalalignment='top')
-            axs[i].text(0.02, 0.95, text_str, **text_kwargs)
-
     # Scatter plot
     if do_scatter_plot:
         i=1
-        default_sc_kwargs = dict(alpha=0.6, ls='none')
-        sc_kwargs = {**default_sc_kwargs, **sc_kwargs}
-
         axs[i].set_ylabel("Antenna no.")
         axs[i].plot(measured_phases, np.arange(len(measured_phases)), marker='x' if plot_residuals else '3', color=colors[0], label='Measured', **sc_kwargs)
 
@@ -92,6 +95,9 @@ def phase_comparison_figure(
 
     fig.tight_layout()
 
+    if return_fit_info:
+        return fig, fit_info
+
     return fig
 
 
@@ -210,4 +216,4 @@ def fitted_histogram_figure(
         loc = (0.98, 0.95)
         ax.text(*loc, text_str, **{**text_kwargs, **dict(ha='right')})
 
-    return fig, fit_info
+    return fig, dict(fit_info=fit_info, counts=counts, bins=bins)