mirror of
https://gitlab.science.ru.nl/mthesis-edeboone/m-thesis-introduction.git
synced 2024-11-13 10:03:32 +01:00
428 lines
13 KiB
Python
428 lines
13 KiB
Python
#!/usr/bin/env python3
|
|
|
|
__doc__ = \
|
|
"""
|
|
Show the curve for signal-to-noise ratio vs N_samples
|
|
"""
|
|
|
|
from collections import namedtuple
|
|
|
|
import matplotlib.pyplot as plt
|
|
import matplotlib.gridspec as gridspec
|
|
import numpy as np
|
|
import scipy.fftpack as ft
|
|
|
|
rng = np.random.default_rng()
|
|
|
|
passband = namedtuple("Band", ['low', 'high'], defaults=[0, np.inf])
|
|
|
|
def get_freq_spec(val,dt):
|
|
"""From earsim/tools.py"""
|
|
fval = np.fft.fft(val)[:len(val)//2]
|
|
freq = np.fft.fftfreq(len(val),dt)[:len(val)//2]
|
|
return fval, freq
|
|
|
|
|
|
def ft_spectrum( signal, sample_rate=1, ftfunc=None, freqfunc=None, mask_bias=False, normalise_amplitude=False):
|
|
"""Return a FT of $signal$, with corresponding frequencies"""
|
|
|
|
if True:
|
|
return get_freq_spec(signal, 1/sample_rate)
|
|
|
|
n_samples = len(signal)
|
|
|
|
if ftfunc is None:
|
|
real_signal = np.isrealobj(signal)
|
|
if False and real_signal:
|
|
ftfunc = ft.rfft
|
|
freqfunc = ft.rfftfreq
|
|
else:
|
|
ftfunc = ft.fft
|
|
freqfunc = ft.fftfreq
|
|
|
|
if freqfunc is None:
|
|
freqfunc = ft.fftfreq
|
|
|
|
normalisation = 2/len(signal) if normalise_amplitude else 1
|
|
|
|
spectrum = normalisation * ftfunc(signal)
|
|
freqs = freqfunc(n_samples, 1/sample_rate)
|
|
|
|
if not mask_bias:
|
|
return spectrum, freqs
|
|
else:
|
|
return spectrum[1:], freqs[1:]
|
|
|
|
|
|
def plot_spectrum( spectrum, freqs, plot_complex=False, plot_power=False, plot_amplitude=None, ax=None, freq_unit="Hz", freq_scaler=1):
|
|
""" Plot a signal's spectrum on an Axis object"""
|
|
plot_amplitude = plot_amplitude or (not plot_power and not plot_complex)
|
|
alpha = 1
|
|
|
|
if ax is None:
|
|
ax = plt.gca()
|
|
|
|
ax.set_title("Spectrum")
|
|
ax.set_xlabel("f" + (" ["+freq_unit+"]" if freq_unit else "" ))
|
|
ylabel = ""
|
|
if plot_amplitude or plot_complex:
|
|
ylabel = "Amplitude"
|
|
if plot_power:
|
|
if ylabel:
|
|
ylabel += "|"
|
|
ylabel += "Power"
|
|
ax.set_ylabel(ylabel)
|
|
|
|
if plot_complex:
|
|
alpha = 0.5
|
|
ax.plot(freqs/freq_scaler, np.real(spectrum), '.-', label='Real', alpha=alpha)
|
|
ax.plot(freqs/freq_scaler, np.imag(spectrum), '.-', label='Imag', alpha=alpha)
|
|
|
|
if plot_power:
|
|
ax.plot(freqs/freq_scaler, np.abs(spectrum)**2, '.-', label='Power', alpha=alpha)
|
|
|
|
if plot_amplitude:
|
|
ax.plot(freqs/freq_scaler, np.abs(spectrum), '.-', label='Abs', alpha=alpha)
|
|
|
|
ax.legend()
|
|
|
|
return ax
|
|
|
|
def plot_phase( spectrum, freqs, ylim_epsilon=0.5, ax=None, freq_unit="Hz", freq_scaler=1):
|
|
if ax is None:
|
|
ax = plt.gca()
|
|
|
|
ax.set_ylabel("Phase")
|
|
ax.set_xlabel("f" + (" ["+freq_unit+"]" if freq_unit else "" ))
|
|
|
|
ax.plot(freqs/freq_scaler, np.angle(spectrum), '.-')
|
|
ax.set_ylim(-1*np.pi - ylim_epsilon, np.pi + ylim_epsilon)
|
|
|
|
return ax
|
|
|
|
def plot_signal( signal, sample_rate = 1, ax=None, time=None, time_unit="s", **kwargs):
|
|
if ax is None:
|
|
ax = plt.gca()
|
|
|
|
if time is None:
|
|
time = np.arange(len(signal))/sample_rate
|
|
|
|
ax.set_title("Signal")
|
|
ax.set_xlabel("t" + (" ["+time_unit+"]" if time_unit else "" ))
|
|
ax.set_ylabel("A(t)")
|
|
|
|
ax.plot(time, signal, **kwargs)
|
|
|
|
return ax
|
|
|
|
def plot_combined_spectrum(spectrum, freqs,
|
|
spectrum_kwargs={}, fig=None, gs=None, freq_scaler=1, freq_unit="Hz"):
|
|
"""Plot both the frequencies and phase in one figure."""
|
|
|
|
# configure plotting layout
|
|
if fig is None:
|
|
fig = plt.figure(figsize=(8, 16))
|
|
|
|
if gs is None:
|
|
gs = gridspec.GridSpec(2, 1, figure=fig, height_ratios=[3,1], hspace=0)
|
|
|
|
ax1 = fig.add_subplot(gs[:-1, -1])
|
|
ax2 = fig.add_subplot(gs[-1, -1], sharex=ax1)
|
|
|
|
axes = np.array([ax1, ax2])
|
|
|
|
# plot the spectrum
|
|
plot_spectrum(spectrum, freqs, ax=ax1, freq_scaler=freq_scaler, freq_unit=freq_unit, **spectrum_kwargs)
|
|
|
|
# plot the phase
|
|
plot_phase(spectrum, freqs, ax=ax2, freq_scaler=freq_scaler, freq_unit=freq_unit)
|
|
|
|
ax1.xaxis.tick_top()
|
|
[label.set_visible(False) for label in ax1.get_xticklabels()]
|
|
|
|
return fig, axes
|
|
|
|
|
|
def phasemod(phase, low=np.pi):
|
|
"""
|
|
Modulo phase such that it falls within the
|
|
interval $[-low, 2\pi - low)$.
|
|
"""
|
|
return (phase + low) % (2*np.pi) - low
|
|
|
|
def save_all_figs_to_path(fnames, figs=None, default_basename=__file__, default_extensions=['.pdf', '.png']):
|
|
if figs is None:
|
|
figs = [plt.figure(i) for i in plt.get_fignums()]
|
|
|
|
default_basename = path.basename(default_basename)
|
|
|
|
# singular value
|
|
if isinstance(fnames, (str, True)):
|
|
fnames = [fnames]
|
|
|
|
if len(fnames) == len(figs):
|
|
fnames_list = zip(figs, fnames, False)
|
|
elif len(fnames) == 1:
|
|
tmp_fname = fnames[0] #needed for generator
|
|
fnames_list = ( (fig, tmp_fname, len(figs) > 1) for fig in figs)
|
|
else:
|
|
# outer product magic
|
|
fnames_list = ( (fig,fname, False) for fname in fnames for fig in figs )
|
|
del fnames
|
|
# format fnames
|
|
pad_width = max(2, int(np.floor(np.log10(len(figs))+1)))
|
|
|
|
fig_fnames = []
|
|
for fig, fnames, append_num in fnames_list:
|
|
if not hasattr(fnames, '__len__') or isinstance(fnames, str):
|
|
# single name
|
|
fnames = [fnames]
|
|
|
|
new_fnames = []
|
|
for fname in fnames:
|
|
if path.isdir(fname):
|
|
fname = path.join(fname, path.splitext(default_basename)[0]) # leave off extension
|
|
if append_num is True:
|
|
fname += ("_fig{:0"+str(pad_width)+"d}").format(fig.number)
|
|
|
|
if not path.splitext(fname)[1]: # no extension
|
|
for ext in default_extensions:
|
|
new_fnames.append(fname+ext)
|
|
else:
|
|
new_fnames.append(fname)
|
|
|
|
fig_fnames.append(new_fnames)
|
|
|
|
# save files
|
|
for fnames, fig in zip(fig_fnames, figs):
|
|
for fname in fnames:
|
|
fig.savefig(fname, transparent=True)
|
|
|
|
def sine_fitfunc(t, amp=1, freq=1, phase=0, off=0):
|
|
"""Simple sine wave for fitting purposes"""
|
|
return amp*np.sin( 2*np.pi*freq*t + phase) + off
|
|
|
|
def sampled_time(sample_rate=1, start=0, end=1, offset=0):
|
|
return offset + np.arange(start, end, 1/sample_rate)
|
|
|
|
|
|
def bandpass_mask(freqs, band=passband()):
|
|
low_pass = abs(freqs) <= band[1]
|
|
high_pass = abs(freqs) >= band[0]
|
|
|
|
return low_pass & high_pass
|
|
|
|
def bandsize(band = passband()):
|
|
return band[1] - band[0]
|
|
|
|
def bandlevel(samples, samplerate=1, band=passband(), normalise_bandsize=True, **ft_kwargs):
|
|
fft, freqs = ft_spectrum(samples, samplerate, **ft_kwargs)
|
|
|
|
bandmask = bandpass_mask(freqs, band=band)
|
|
|
|
if normalise_bandsize:
|
|
bins = np.count_nonzero(bandmask, axis=-1)
|
|
else:
|
|
bins = 1
|
|
|
|
level = np.sum(np.abs(fft[bandmask])**2)
|
|
|
|
return level/bins
|
|
|
|
|
|
def noisy_sine_sampling(time, init_params, noise_sigma=1, rng=rng):
|
|
if init_params[2] is None:
|
|
init_params[2] = phasemod(2*np.pi*rng.random())
|
|
|
|
samples = sine_fitfunc(time, *init_params)
|
|
noise = rng.normal(0, noise_sigma, size=len(samples))
|
|
|
|
|
|
return samples, noise
|
|
|
|
def main(
|
|
N = 1,
|
|
f_sample = 250e6, # Hz
|
|
t_length = 1e4 * 1e-9, # s
|
|
|
|
noise_band = passband(30e6, 80e6),
|
|
noise_sigma = 1,
|
|
|
|
# signal properties
|
|
f_sine = 50e6,
|
|
signal_band = passband(50e6 - 1e6, 50e6 + 1e6),
|
|
sine_amp = 0.2,
|
|
sine_offset = 0,
|
|
return_ranges_plot = False,
|
|
cut_signal_band_from_noise_band = False
|
|
):
|
|
N = int(N)
|
|
|
|
init_params = np.array([sine_amp, f_sine, None, sine_offset])
|
|
|
|
axs = None
|
|
snrs = np.zeros( N )
|
|
time = sampled_time(f_sample, end=t_length)
|
|
for j in range(N):
|
|
samples, noise = noisy_sine_sampling(time, init_params, noise_sigma)
|
|
|
|
|
|
# determine signal to noise
|
|
noise_level = bandlevel(noise, f_sample, noise_band)
|
|
if cut_signal_band_from_noise_band:
|
|
lower_noise_band = passband(noise_band[0], signal_band[0])
|
|
upper_noise_band = passband(signal_band[1], noise_band[1])
|
|
|
|
noise_level = bandlevel(noise, f_sample, lower_noise_band)
|
|
noise_level += bandlevel(noise, f_sample, upper_noise_band)
|
|
|
|
signal_level = bandlevel(samples, f_sample, signal_band)
|
|
|
|
snrs[j] = np.sqrt(signal_level/noise_level)
|
|
|
|
# make a nice plot showing what ranges were taken
|
|
# and the bandlevels associated with them
|
|
if return_ranges_plot and j == 0:
|
|
combined_fft, freqs = ft_spectrum(samples+noise, f_sample)
|
|
|
|
# plot the original signal
|
|
if False:
|
|
_, ax = plt.subplots()
|
|
ax = plot_signal(samples+noise, sample_rate=f_sample/1e6, time_unit='us', ax=ax)
|
|
|
|
# plot the spectrum
|
|
if True:
|
|
freq_scaler=1e6
|
|
_, axs = plot_combined_spectrum(combined_fft, freqs, freq_scaler=freq_scaler, freq_unit='MHz')
|
|
|
|
# indicate band ranges and frequency
|
|
for ax in axs:
|
|
ax.axvline(f_sine/freq_scaler, color='r', alpha=0.4)
|
|
ax.axvspan(noise_band[0]/freq_scaler, noise_band[1]/freq_scaler, color='purple', alpha=0.3, label='noiseband')
|
|
ax.axvspan(signal_band[0]/freq_scaler, signal_band[1]/freq_scaler, color='orange', alpha=0.3, label='signalband')
|
|
|
|
# indicate initial phase
|
|
axs[1].axhline(init_params[2], color='r', alpha=0.4)
|
|
|
|
# plot the band levels
|
|
levelax = axs[0].twinx()
|
|
levelax.set_ylabel("Bandlevel")
|
|
levelax.hlines(signal_level, noise_band[0]/freq_scaler, signal_band[1]/freq_scaler, colors=['orange'])
|
|
levelax.hlines(noise_level, noise_band[0]/freq_scaler, noise_band[1]/freq_scaler, colors=['purple'])
|
|
levelax.set_ylim(bottom=0)
|
|
|
|
axs[0].legend()
|
|
|
|
# plot signal_band pass signal
|
|
if False:
|
|
freqs = np.fft.fftfreq(len(samples), 1/f_sample)
|
|
bandmask = bandpass_mask(freqs, band=signal_band)
|
|
fft = np.fft.fft(samples)
|
|
fft[ ~bandmask ] = 0
|
|
bandpassed_samples = np.fft.ifft(fft)
|
|
|
|
_, ax3 = plt.subplots()
|
|
ax3 = plot_signal(bandpassed_samples, sample_rate=f_sample/1e6, time_unit='us', ax=ax3)
|
|
ax3.set_title("Bandpassed Signal")
|
|
|
|
|
|
return snrs, axs
|
|
|
|
|
|
if __name__ == "__main__":
|
|
from argparse import ArgumentParser
|
|
import os.path as path
|
|
|
|
rng = np.random.default_rng(1)
|
|
|
|
parser = ArgumentParser(description=__doc__)
|
|
parser.add_argument("fname", metavar="path/to/figure[/]", nargs="?", help="Location for generated figure, will append __file__ if a directory. If not supplied, figure is shown.")
|
|
|
|
args = parser.parse_args()
|
|
default_extensions = ['.pdf', '.png']
|
|
|
|
if args.fname == 'none':
|
|
args.fname = None
|
|
|
|
###
|
|
t_lengths = np.linspace(1e3, 5e4)* 1e-9 # s
|
|
N = 10e1
|
|
f_sine = 53.3e6 # Hz
|
|
f_sample = 250e6 # Hz
|
|
|
|
if False:
|
|
N = 1 # Note: keep this low, N figures will be displayed!
|
|
N_t_length = 10
|
|
for t_length in t_lengths[-N_t_length-1:-1]:
|
|
snrs = np.zeros( int(N))
|
|
for i in range(int(N)):
|
|
delta_f = 1/t_length
|
|
snrs[i], axs = main(
|
|
N=1,
|
|
t_length=t_length,
|
|
f_sample=f_sample,
|
|
|
|
# signal properties
|
|
f_sine = f_sine,
|
|
sine_amp = 1,
|
|
noise_sigma = 1,
|
|
|
|
noise_band = passband(30e6, 80e6),
|
|
signal_band = passband(f_sine- 3*delta_f, f_sine + 3*delta_f),
|
|
|
|
return_ranges_plot=True
|
|
)
|
|
|
|
axs[0].set_title("SNR: {}, N:{}".format(snrs[i], t_length*f_sample))
|
|
axs[0].set_xlim(
|
|
(f_sine - 20*delta_f)/1e6,
|
|
(f_sine + 20*delta_f)/1e6
|
|
)
|
|
|
|
print(snrs, "M:",np.mean(snrs))
|
|
|
|
plt.show(block=False)
|
|
|
|
else:
|
|
#original code
|
|
my_snrs = np.zeros( (len(t_lengths), int(N)) )
|
|
for j, t_length in enumerate(t_lengths):
|
|
return_ranges_plot = ((j==0) and True) or ( (j==(len(t_lengths)-1)) and True)
|
|
|
|
delta_f = 1/t_length
|
|
|
|
my_snrs[j], axs = main(
|
|
N=N,
|
|
t_length=t_length,
|
|
f_sample = f_sample,
|
|
|
|
# signal properties
|
|
f_sine = f_sine,
|
|
sine_amp = 1,
|
|
noise_sigma = 1,
|
|
|
|
noise_band = passband(30e6, 80e6),
|
|
signal_band = passband(f_sine- 3*delta_f, f_sine + 3*delta_f),
|
|
|
|
return_ranges_plot=return_ranges_plot,
|
|
)
|
|
|
|
if return_ranges_plot:
|
|
ranges_axs = axs
|
|
|
|
fig, axs2 = plt.subplots()
|
|
axs2.set_xlabel("N = T*$f_s$")
|
|
axs2.set_ylabel("SNR")
|
|
|
|
for j, t_length in enumerate(t_lengths):
|
|
t_length = t_length * f_sample
|
|
axs2.plot(np.repeat(t_length, my_snrs.shape[1]), my_snrs[j], ls='none', color='blue', marker='o', alpha=max(0.01, 1/my_snrs.shape[1]))
|
|
# plot the means
|
|
axs2.plot(t_lengths*f_sample, np.mean(my_snrs, axis=-1), color='green', marker='*', ls='none')
|
|
|
|
### Save or show figures
|
|
if not args.fname:
|
|
# empty list, False, None
|
|
plt.show()
|
|
else:
|
|
save_all_figs_to_path(args.fname, default_basename=__file__)
|