m-thesis-introduction/simulations/11_pulsed_timing.py

#!/usr/bin/env python3


from lib import util
from scipy import signal, interpolate, stats
import matplotlib.pyplot as plt
import numpy as np
from itertools import zip_longest
import h5py
from copy import deepcopy

try:
    from tqdm import tqdm
except:
    tqdm = lambda x: x

rng = np.random.default_rng()

class Waveform:
    name = None
    signal = None
    dt = None

    _t = None

    def __init__(self,signal=None, dt=None, t=None, name=None):
        self.signal = signal
        self.name = name

        if t is not None:
            assert len(t) == len(signal)

            self._t = t
            self.dt = t[1] - t[0]
        elif dt is not None:
            self.dt = dt

    # Lazy evaluation of time
    @property
    def t(self):
        if self._t is None:
            return self.dt * np.arange(0, len(self.signal))
        return self._t

    @t.setter
    def t(self, value):
        self._t = value

    @t.deleter
    def t(self):
        del self._t

    def __len__():
        return len(self.signal)

def white_noise_realisation(N_samples, noise_sigma=1, rng=rng):
    return rng.normal(0, noise_sigma or 0, size=N_samples)

def antenna_bp(trace, low_bp, high_bp, dt, order=3):
    fs = 1/dt

    bp_filter = signal.butter(order, [low_bp, high_bp], 'band', fs=fs, output='sos')
    bandpassed = signal.sosfilt(bp_filter, trace)

    return bandpassed

def my_correlation(in1, template, lags=None):
    template_length = len(template)
    in1_length = len(in1)

    if lags is None:
        lags = np.arange(-template_length+1, in1_length + 1)

    # do the correlation jig
    corrs = np.zeros_like(lags, dtype=float)
    for i, l in enumerate(lags):
        if l <= 0: # shorten template at the front
            in1_start = 0
            template_end = template_length

            template_start = -template_length - l
            in1_end = max(0, min(in1_length, -template_start)) # 0 =< l + template_length =< in1_lengt

        elif l > in1_length - template_length:
            # shorten template from the back
            in1_end = in1_length
            template_start = 0

            in1_start = min(l, in1_length)
            template_end = max(0,  in1_length - l)

        else:
            in1_start = min(l, in1_length)
            in1_end   = min(in1_start + template_length, in1_length)

            # full template
            template_start = 0
            template_end   = template_length

        # Slice in1 and template
        in1_slice = in1[in1_start:in1_end]
        template_slice = template[template_start:template_end]

        corrs[i] = np.dot(in1_slice, template_slice)

    return corrs, (in1, template, lags)

def trace_upsampler(trace, template_t, trace_t):
    template_dt = template.t[1] - template.t[0]
    trace_dt = trace_t[1] - trace_t[0]

    upsample_factor = trace_dt/template_dt
    upsampled_trace_N = np.ceil(len(trace) * upsample_factor)

    upsample_factor = int(upsample_factor)
    upsampled_trace_N = int(upsampled_trace_N)

    # upsample trace
    upsampled_trace = np.zeros(upsampled_trace_N)
    upsampled_trace[::upsample_factor] = trace

    #upsampled_t = np.arange(trace_t[0], trace_t[-1], template_dt)
    upsampled_t = template_dt * np.arange(len(upsampled_trace)) + trace_t[0]

    return upsampled_trace, upsampled_t

def trace_downsampler(trace, template_t, trace_t, offset):
    pass

def read_time_residuals_cache(cache_fname, template_dt, antenna_dt, snr_sigma_factor, N=None):
    try:
        with h5py.File(cache_fname, 'r') as fp:
            pgroup = fp['time_residuals']
            pgroup2 = pgroup[f'{template_dt}_{antenna_dt}']

            ds_name = str(snr_sigma_factor)

            ds = pgroup2[ds_name]
            if N is None:
                return deepcopy(ds[:])
            else:
                return deepcopy(ds[:min(N, len(ds))])
    except (KeyError, FileNotFoundError):
        return np.array([])

def write_time_residuals_cache(cache_fname, time_residuals, template_dt, antenna_dt, noise_sigma_factor):
    with h5py.File(cache_fname, 'a') as fp:
        pgroup = fp.require_group('time_residuals')
        pgroup2 = pgroup.require_group(f'{template_dt}_{antenna_dt}')

        ds_name = str(noise_sigma_factor)

        if ds_name in pgroup2.keys():
            del pgroup2[ds_name]

        ds = pgroup2.create_dataset(ds_name, (len(time_residuals)), dtype='f', data=time_residuals, maxshape=(None))

if __name__ == "__main__":
    import os
    import matplotlib
    import sys
    if os.name == 'posix' and "DISPLAY" not in os.environ:
        matplotlib.use('Agg')

    bp_freq = (30e-3, 80e-3) # GHz
    template_dt = 5e-2 # ns
    interp_template_dt = 5e-5 # ns
    template_length = 200 # ns

    N_residuals = 50*3 if len(sys.argv) < 2 else int(sys.argv[1])
    snr_factors = np.concatenate( # 1/noise_amplitude factor
            (
                [0.25, 0.5, 0.75],
                [1, 1.5, 2, 2.5, 3, 4, 5, 7],
                [10, 20, 30, 50],
                [100, 200, 300, 500]
            ),
            axis=None)

    antenna_dt = 2 # ns
    antenna_timelength = 1024 # ns

    cut_wrong_peak_matches = True

    #
    # Create the template
    #
    template = Waveform(None, dt=template_dt, name='Template')
    _deltapeak = util.deltapeak(timelength=template_length, samplerate=1/template.dt, offset=0)
    template.signal = antenna_bp(_deltapeak[0], *bp_freq, template_dt)
    template.peak_sample = _deltapeak[1]
    template.peak_time = template.dt * template.peak_sample

    # Interpolation Template
    # to create an 'analog' sampled antenna
    interp_template = Waveform(None, dt=interp_template_dt, name='Interpolation Template')
    _interp_deltapeak = util.deltapeak(timelength=template_length, samplerate=1/interp_template.dt, offset=0)
    interp_template.signal = antenna_bp(_interp_deltapeak[0], *bp_freq, interp_template.dt)
    interp_template.peak_sample = _interp_deltapeak[1]
    interp_template.peak_time = interp_template.dt * interp_template.peak_sample
    interp_template.signal *= max(template.signal)/max(interp_template.signal)
    interp_template.interpolate = interpolate.interp1d(interp_template.t, interp_template.signal, assume_sorted=True, fill_value=0, bounds_error=False)

    if True: # show template
        fig, ax = plt.subplots()
        ax.set_title("Deltapeak and Bandpassed Template")
        ax.set_xlabel("Time [ns]")
        ax.set_ylabel("Amplitude")
        ax.plot(template.t, max(template.signal)*_deltapeak[0], label='Impulse Template')
        ax.plot(template.t, template.signal, label='Filtered Template')
        ax.plot(interp_template.t, interp_template.signal, label='Filtered Interpolation Template', ls='dashed')
        ax.legend()
        fig.savefig('figures/11_template_deltapeak.pdf')

        if True: # show filtering equivalence samplerates
            _deltapeak = util.deltapeak(timelength=template_length, samplerate=1/antenna_dt, offset=0)
            _time = util.sampled_time(end=template_length, sample_rate=1/antenna_dt)
            _bandpassed = antenna_bp(_deltapeak[0], *bp_freq, antenna_dt)

            ax.plot(_time, max(_bandpassed)*_deltapeak[0], label='Impulse Antenna')
            ax.plot(_time, _bandpassed, label='Filtered Antenna')

            ax.legend()
            fig.savefig('figures/11_template_deltapeak+antenna.pdf')

        if True:
            plt.close(fig)

    #
    # Find time accuracies as a function of signal strength
    #
    h5_cache_fname = f'11_pulsed_timing.hdf5'

    time_accuracies = np.zeros(len(snr_factors))
    for k, snr_sigma_factor in tqdm(enumerate(snr_factors)):
        # Read in cached time residuals
        if True:
            cached_time_residuals = read_time_residuals_cache(h5_cache_fname, template.dt, antenna_dt, snr_sigma_factor)
        else:
            cached_time_residuals = np.array([])

        #
        # Find difference between true and templated times
        #
        time_residuals = np.zeros(max(0, (N_residuals - len(cached_time_residuals))))
        for j in tqdm(range(len(time_residuals))):
            do_plots = j==0

            # receive at antenna
            ## place the deltapeak signal at a random location
            antenna = Waveform(None, dt=antenna_dt, name='Signal')

            if False: # Create antenna trace without template
                antenna_true_signal, antenna_peak_sample = util.deltapeak(timelength=antenna_timelength, samplerate=1/antenna.dt, offset=[0.2, 0.8], rng=rng)

                antenna.peak_sample = antenna_peak_sample
                antenna.peak_time = antenna.dt * antenna.peak_sample
                antenna.signal = antenna_bp(antenna.signal, *bp_freq, antenna.dt)
                print(f"Antenna Peak Time: {antenna.peak_time}")
                print(f"Antenna Peak Sample: {antenna.peak_sample}")

            else: # Sample the template at some offset
                antenna.peak_time = antenna_timelength * ((0.8 - 0.2) *rng.random(1) + 0.2)
                sampling_offset = rng.random(1)*antenna.dt

                antenna.t = util.sampled_time(1/antenna.dt, start=0, end=antenna_timelength)

                # Sample the interpolation template
                antenna.signal = interp_template.interpolate(antenna.t - antenna.peak_time)

                antenna.peak_sample = antenna.peak_time/antenna.dt
                antenna_true_signal = antenna.signal

            true_time_offset = antenna.peak_time - template.peak_time

            if False: # flip polarisation
                antenna.signal *= -1

            ## Add noise
            noise_amplitude = max(template.signal) * 1/snr_sigma_factor
            noise_realisation = noise_amplitude * white_noise_realisation(len(antenna.signal))
            filtered_noise = antenna_bp(noise_realisation, *bp_freq, antenna.dt)

            antenna.signal += filtered_noise

            if do_plots: # show signals
                fig, axs = plt.subplots(2, sharex=True)
                axs[0].set_title("Antenna Waveform")
                axs[-1].set_xlabel("Time [ns]")
                axs[0].set_ylabel("Amplitude")
                axs[0].plot(antenna.t, antenna.signal, label='bandpassed w/ noise', alpha=0.9)
                axs[0].plot(antenna.t, antenna.signal - filtered_noise, label='bandpassed w/o noise', alpha=0.9)
                axs[0].legend()

                axs[1].set_title("Template")
                axs[1].set_ylabel("Amplitude")
                axs[1].plot(template.t, template.signal, label='orig')
                axs[1].plot(template.t + true_time_offset, template.signal, label='true moved orig')
                axs[1].legend()

                axs[0].grid()
                axs[1].grid()

                fig.savefig('figures/11_antenna_signals.pdf')

                if True: # zoom
                    wx = 100
                    x0 = true_time_offset

                    old_xlims = axs[0].get_xlim()
                    axs[0].set_xlim( x0-wx, x0+wx)
                    fig.savefig('figures/11_antenna_signals_zoom.pdf')

                    # restore
                    axs[0].set_xlim(*old_xlims)

                if True:
                    plt.close(fig)

            axs2 = None
            if True: # upsampled trace
                upsampled_trace, upsampled_t = trace_upsampler(antenna.signal, template.t, antenna.t)

                if do_plots: # Show upsampled traces
                    fig2, axs2 = plt.subplots(1, sharex=True)
                    if not hasattr(axs2, '__len__'):
                        axs2 = [axs2]

                    axs2[-1].set_xlabel("Time [ns]")
                    axs2[0].set_ylabel("Amplitude")
                    axs2[0].plot(antenna.t, antenna.signal, marker='o', label='orig')
                    axs2[0].plot(upsampled_t, upsampled_trace, label='upsampled')
                    axs2[0].legend(loc='upper right')

                    fig2.savefig('figures/11_upsampled.pdf')

                    wx = 1e2
                    x0 = upsampled_t[0] + wx - 5
                    axs2[0].set_xlim(x0-wx, x0+wx)
                    fig2.savefig('figures/11_upsampled_zoom.pdf')

                    if True:
                        plt.close(fig2)

                # determine correlations with arguments
                lag_dt = upsampled_t[1] - upsampled_t[0]
                corrs, (out1_signal, out2_template, lags) = my_correlation(upsampled_trace, template.signal)

                # Determine best correlation time
                idx = np.argmax(abs(corrs))
                best_sample_lag = lags[idx]
                best_time_lag = best_sample_lag * lag_dt

            else: # downsampled template
                raise NotImplementedError

                corrs, (_, _, lags) = my_downsampling_correlation(antenna.signal, antenna.t, template.signal, template.t)
                lag_dt = upsampled_t[1] - upsampled_t[0]

            # Calculate the time residual
            time_residuals[j] = best_time_lag - true_time_offset

            if not do_plots:
                continue

            if do_plots and axs2:
                axs2[-1].axvline(best_time_lag, color='r', alpha=0.5, linewidth=2)
                axs2[-1].axvline(true_time_offset, color='g', alpha=0.5, linewidth=2)

            # Show the final signals correlated
            if do_plots:
                # amplitude scaling required for single axis plotting
                template_amp_scaler = max(abs(template.signal)) / max(abs(antenna.signal))

                # start the figure
                fig, axs = plt.subplots(2, sharex=True)
                ylabel_kwargs = dict(
                                    #rotation=0,
                                    ha='right',
                                    va='center'
                                )
                axs[-1].set_xlabel("Time [ns]")

                offset_list = [
                        [best_time_lag, dict(label=template.name, color='orange')],
                        [true_time_offset, dict(label='True offset', color='green')],
                        ]

                # Signal
                i=0
                axs[i].set_ylabel("Amplitude", **ylabel_kwargs)
                axs[i].plot(antenna.t, antenna.signal, label=antenna.name)

                # Plot the template
                for offset_args in offset_list:
                    this_kwargs = offset_args[1]
                    offset = offset_args[0]

                    l = axs[i].plot(offset + template.t, template_amp_scaler * template.signal, **this_kwargs)

                axs[i].legend()

                # Correlation
                i=1
                axs[i].set_ylabel("Correlation", **ylabel_kwargs)
                axs[i].plot(lags * lag_dt, corrs)

                # Lines across both axes
                for offset_args in offset_list:
                    this_kwargs = offset_args[1]
                    offset = offset_args[0]

                    for i in [0,1]:
                        axs[i].axvline(offset, ls='--', color=this_kwargs['color'], alpha=0.7)

                    axs[0].axvline(offset + len(template.signal) * (template.t[1] - template.t[0]), color=this_kwargs['color'], alpha=0.7)


                if True: # zoom
                    wx = len(template.signal) * (template.dt)/2
                    t0 = best_time_lag

                    old_xlims = axs[0].get_xlim()
                    axs[i].set_xlim( x0-wx, x0+3*wx)
                    fig.savefig('figures/11_corrs_zoom.pdf')

                    # restore
                    axs[i].set_xlim(*old_xlims)

                fig.tight_layout()
                fig.savefig('figures/11_corrs.pdf')

                if True:
                    plt.close(fig)

        print()# separating tqdm
        print()# separating tqdm
        # Were new time residuals calculated?
        # Add them to the cache file
        if len(time_residuals) > 1:
            # merge cached and calculated time residuals
            time_residuals = np.concatenate((cached_time_residuals, time_residuals), axis=None)

            if True: # write the cache
                write_time_residuals_cache(h5_cache_fname, time_residuals, template_dt, antenna_dt, snr_sigma_factor)
        else:
            time_residuals = cached_time_residuals

        # Make a plot of the time residuals
        if N_residuals > 1:
            time_accuracies[k] = np.std(time_residuals[:N_residuals])

            hist_kwargs = dict(bins='sqrt', density=False, alpha=0.8, histtype='step')
            fig, ax = plt.subplots()
            ax.set_title(
                    "Template Correlation Lag finding"
                    + f"\n template dt: {template_dt*1e3: .1e}ps"
                    + f"; antenna dt: {antenna_dt: .1e}ns"
                    + f"; noise_factor: {noise_sigma_factor: .1e}"
                    )
            ax.set_xlabel("Time Residual [ns]")
            ax.set_ylabel("#")

            counts, bins, _patches = ax.hist(time_residuals, **hist_kwargs)
            if True: # fit gaussian to histogram
                min_x = min(time_residuals)
                max_x = max(time_residuals)

                dx = bins[1] - bins[0]
                scale = len(time_residuals) * dx

                xs = np.linspace(min_x, max_x)

                # do the fit
                name = "Norm"
                param_names = [ "$\\mu$", "$\\sigma$" ]
                distr_func = stats.norm

                label = name +"(" + ','.join(param_names) + ')'

                # plot
                fit_params = distr_func.fit(time_residuals)
                fit_ys = scale * distr_func.pdf(xs, *fit_params)
                ax.plot(xs, fit_ys, label=label)

                # chisq
                ct = np.diff(distr_func.cdf(bins, *fit_params))*np.sum(counts)
                if True:
                    ct *= np.sum(counts)/np.sum(ct)
                c2t = stats.chisquare(counts, ct, ddof=len(fit_params))
                chisq_strs = [
                            f"$\\chi^2$/dof = {c2t[0]: .2g}/{len(fit_params)}"
                            ]

                # text on plot
                text_str = "\n".join(
                    [label]
                    +
                    [ f"{param} = {value: .2e}" for param, value in zip_longest(param_names, fit_params, fillvalue='?') ]
                    +
                    chisq_strs
                    )

                ax.text( *(0.02, 0.95), text_str, fontsize=12, ha='left', va='top', transform=ax.transAxes)

            fig.savefig(f"figures/11_time_residual_hist_tdt{template_dt:0.1e}_n{noise_sigma_factor: .1e}.pdf")

            if True:
                plt.close(fig)

    # SNR time accuracy plot
    if True:
        fig, ax = plt.subplots()
        ax.set_title(f"Template matching SNR vs time accuracy")
        ax.set_xlabel("Signal to Noise Factor")
        ax.set_ylabel("Time Accuracy [ns]")

        ax.legend(title="\n".join([
            f"N={N_residuals}",
            f"template_dt={template_dt:0.1e}ns",
            f"antenna_dt={antenna_dt:0.1e}ns",
            ]))

        if True:
            ax.set_xscale('log')
            ax.set_yscale('log')

        # plot the values
        ax.plot(np.asarray(snr_factors), time_accuracies, ls='none', marker='o')

        if True: # limit y-axis to 1e0
            ax.set_ylim([None, 1e1])


        # Set horizontal line at 1 ns
        if True:
            ax.axhline(1, ls='--', alpha=0.8, color='g')
            ax.grid()
            ax.axhline(template_dt/np.sqrt(12), ls='--', alpha=0.7, color='b')

        fig.tight_layout()
        fig.savefig(f"figures/11_time_res_vs_snr_tdt{template_dt:0.1e}.pdf")

    plt.show()