ZH: plot noise parameters if available

+ figlib: histogram with distr fitting

simulations/airshower_beacon_simulation/ba_measure_beacon_phase.py

@@ -12,6 +12,8 @@ import numpy as np
 
 import aa_generate_beacon as beacon
 import lib
+from lib import figlib
+
 
 if __name__ == "__main__":
     from os import path
@@ -49,6 +51,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
 
@@ -83,7 +86,8 @@ if __name__ == "__main__":
 
         N_antennas = len(group.keys())
         # just for funzies
-        found_data = np.zeros((N_antennas, 3))
+        found_data = np.zeros((N_antennas, 3)) # freq, phase, amp
+        noise_data = np.zeros((N_antennas, 2)) # phase, amp
 
         # Determine frequency and phase
         for i, name in enumerate(group.keys()):
@@ -133,6 +137,7 @@ if __name__ == "__main__":
             # TODO: refine masking
             # use beacon but remove where E_AxB-Beacon != 0
             # Uses the first traces as reference
+            t_mask = 0
             if N_mask and orients[0] != 'B':
                 N_pre, N_post = N_mask//2, N_mask//2
 
@@ -166,6 +171,18 @@ if __name__ == "__main__":
                 beacon_phases = [ 2*np.pi*t0*f_beacon ]
                 amps = [ 3e-7 ]
 
+            # Also try to find the phase from the noise trace if available
+            if len(h5ant[traces_key]) > 4:
+                noise_trace = h5ant[traces_key][5]
+                if np.any(t_mask): # Mask the same area
+                    noise_trace = noise_trace[t_mask]
+
+                real, imag = lib.direct_fourier_transform(f_beacon, t_trace, noise_trace)
+                noise_phase = np.arctan2(imag, real)
+                noise_amp = (real**2 + imag**2)**0.5
+
+                noise_data[i] = noise_phase, noise_amp
+
             # choose highest amp
             idx = np.argmax(amps)
             if False and len(beacon_phases) > 1:
@@ -246,10 +263,16 @@ if __name__ == "__main__":
             h5attrs['amplitude'] = amplitude
             h5attrs['orientation'] = orientation
 
+            if noise_phase:
+                h5attrs['noise_phase'] = noise_phase
+                h5attrs['noise_amp'] = noise_amp
+
     print("Beacon Phases, Amplitudes and Frequencies written to", antennas_fname)
 
     # show histogram of found frequencies
     if show_plots or fig_dir:
+        snrs = beacon.read_snr_file(snr_fname)
+
         if True or allow_frequency_fitting:
             fig, ax = plt.subplots(figsize=figsize)
             ax.set_xlabel("Frequency")
@@ -260,12 +283,45 @@ if __name__ == "__main__":
                 fig.savefig(path.join(fig_dir, path.basename(__file__) + f".hist_freq.pdf"))
 
         if True:
-            fig, ax = plt.subplots(figsize=figsize)
+            fig, _ = figlib.fitted_histogram_figure(found_data[:,2], fit_distr=[], mean_snr=snrs['mean'])
-            ax.set_xlabel("Amplitudes")
+            ax = fig.axes[0]
+            ax.set_xlabel("Amplitude")
             ax.set_ylabel("Counts")
             ax.hist(found_data[:,2], bins='sqrt', density=False)
+            ax.legend()
             if fig_dir:
                 fig.savefig(path.join(fig_dir, path.basename(__file__) + f".hist_amp.pdf"))
 
+        if (noise_data[0] != 0).any():
+            if True:
+                fig, ax = plt.subplots(figsize=figsize)
+                ax.set_title("Noise Phases")
+                ax.set_xlabel("Phase")
+                ax.set_ylabel("#")
+                ax.hist(noise_data[:,0], bins='sqrt', density=False)
+                ax.legend()
+                if fig_dir:
+                    fig.savefig(path.join(fig_dir, path.basename(__file__) + f".noise.hist_phase.pdf"))
+
+            if True:
+                fig, ax = plt.subplots(figsize=figsize)
+                ax.set_title("Noise Phase vs Amplitude")
+                ax.set_xlabel("Phase")
+                ax.set_ylabel("Amplitude [a.u.]")
+                ax.plot(noise_data[:,0], noise_data[:,1], ls='none', marker='x')
+                if fig_dir:
+                    fig.savefig(path.join(fig_dir, path.basename(__file__) + f".noise.phase_vs_amp.pdf"))
+
+            if True:
+                fig, _ = figlib.fitted_histogram_figure(noise_data[:,1], fit_distr=['rice', 'rayleigh'], mean_snr=snrs['mean'])
+                ax = fig.axes[0]
+                ax.set_title("Noise Amplitudes")
+                ax.set_xlabel("Amplitude [a.u.]")
+                ax.set_ylabel("#")
+                ax.legend()
+
+                if fig_dir:
+                    fig.savefig(path.join(fig_dir, path.basename(__file__) + f".noise.hist_amp.pdf"))
+
     if show_plots:
         plt.show()

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