ZH: plot noise parameters if available

+ figlib: histogram with distr fitting

simulations/airshower_beacon_simulation/lib/figlib.py

@@ -1,7 +1,8 @@
 import matplotlib.pyplot as plt
 import numpy as np
 
-from scipy.stats import norm
+import scipy.stats as stats
+from itertools import zip_longest
 
 def phase_comparison_figure(
         measured_phases,
@@ -92,3 +93,106 @@ def phase_comparison_figure(
     fig.tight_layout()
 
     return fig
+
+
+def fitted_histogram_figure(
+        amplitudes,
+        fit_distr = None,
+        text_kwargs={},
+        hist_kwargs={},
+        mean_snr = None,
+        ax = None,
+        **fig_kwargs
+    ):
+    """
+    Create a figure showing $amplitudes$ as a histogram.
+    If fit_distr is a (list of) string, also fit the respective
+    distribution function and show the parameters on the figure.
+    """
+    default_hist_kwargs = dict(bins='sqrt', density=False, alpha=0.8, histtype='step', label='hist')
+    default_text_kwargs = dict(fontsize=14, verticalalignment='top')
+
+    if isinstance(fit_distr, str):
+        fit_distr = [fit_distr]
+
+    hist_kwargs = {**default_hist_kwargs, **hist_kwargs}
+    text_kwargs = {**default_text_kwargs, **text_kwargs}
+
+    if ax is None:
+        fig, ax = plt.subplots(1,1, **fig_kwargs)
+    else:
+        fig = ax.get_figure()
+
+    text_kwargs['transform'] = ax.transAxes
+
+    counts, bins, _patches = ax.hist(amplitudes, **hist_kwargs)
+
+    fit_info = []
+    if fit_distr:
+        min_x = min(amplitudes)
+        max_x = max(amplitudes)
+
+        dx = bins[1] - bins[0]
+        scale = len(amplitudes) * dx
+
+        xs = np.linspace(min_x, max_x)
+
+        for distr in fit_distr:
+            param_manipulate = lambda x: x
+
+            if 'rice' == distr:
+                name = "Rice"
+                param_names = [ "$\\nu$", "$\\sigma$" ]
+                distr_func = stats.rice
+
+                fit_params2text_params = lambda x: (x[0]*x[1], x[1])
+
+            elif 'gaussian' == distr:
+                name = "Norm"
+                param_names = [ "$\\mu$", "$\\sigma$" ]
+                distr_func = stats.norm
+
+            elif 'rayleigh' == distr:
+                name = "Rayleigh"
+                param_names = [ "$\\sigma$" ]
+                distr_func = stats.rayleigh
+
+                fit_params2text_params = lambda x: (x[0]+x[1]/2,)
+
+            else:
+                raise ValueError('Unknown distribution function '+distr)
+
+            label = name +"(" + ','.join(param_names) + ')'
+
+            fit_params = distr_func.fit(amplitudes)
+            fit_ys = distr_func.pdf(xs, *fit_params)
+
+            ax.plot(xs, fit_ys*scale, label=label)
+
+            # change parameters if needed
+            text_fit_params = fit_params2text_params(fit_params)
+
+            text_str = "\n".join(
+                [label]
+                +
+                [ f"{param} = {value: .2e}" for param, value in zip_longest(param_names, text_fit_params, fillvalue='?') ]
+                )
+
+            this_info = {
+                    'name': name,
+                    'param_names': param_names,
+                    'param_values': text_fit_params,
+                    'text_str': text_str,
+                }
+
+            fit_info.append(this_info)
+
+        loc = (0.02, 0.95)
+        ax.text(*loc, "\n\n".join([info['text_str'] for info in fit_info]), **{**text_kwargs, **dict(ha='left')})
+
+    if mean_snr:
+        text_str = f"$\\langle SNR \\rangle$ = {mean_snr: .1e}"
+        loc = (0.98, 0.95)
+        ax.text(*loc, text_str, **{**text_kwargs, **dict(ha='right')})
+
+    return fig, fit_info